Add --cuda-compat-container-root

Signed-off-by: Evan Lezar <elezar@nvidia.com>

Author: elezar

cmd/nvidia-cdi-hook/cudacompat/cudacompat.go

@@ -31,7 +31,7 @@ import (
 )
 
 const (
-	cudaCompatPath = "/usr/local/cuda/compat"
+	defaultCudaCompatPath = "/usr/local/cuda/compat"
 	// cudaCompatLdsoconfdFilenamePattern specifies the pattern for the filename
 	// in ld.so.conf.d that includes a reference to the CUDA compat path.
 	// The 00-compat prefix is chosen to ensure that these libraries have a
@@ -44,8 +44,9 @@ type command struct {
 }
 
 type options struct {
-	hostDriverVersion string
-	containerSpec     string
+	cudaCompatContainerRoot string
+	hostDriverVersion       string
+	containerSpec           string
 }
 
 // NewCommand constructs a cuda-compat command with the specified logger
@@ -76,6 +77,12 @@ func (m command) build() *cli.Command {
 				Usage:       "Specify the host driver version. If the CUDA compat libraries detected in the container do not have a higher MAJOR version, the hook is a no-op.",
 				Destination: &cfg.hostDriverVersion,
 			},
+			&cli.StringFlag{
+				Name:        "cuda-compat-container-root",
+				Usage:       "Specify the folder in which CUDA compat libraries are located in the container",
+				Value:       defaultCudaCompatPath,
+				Destination: &cfg.cudaCompatContainerRoot,
+			},
 			&cli.StringFlag{
 				Name:        "container-spec",
 				Hidden:      true,
@@ -108,7 +115,7 @@ func (m command) run(_ *cli.Command, cfg *options) error {
 		return fmt.Errorf("failed to determined container root: %w", err)
 	}
 
-	containerForwardCompatDir, err := m.getContainerForwardCompatDir(containerRoot(containerRootDir), cfg.hostDriverVersion)
+	containerForwardCompatDir, err := m.getContainerForwardCompatDir(containerRoot(containerRootDir), cfg.cudaCompatContainerRoot, cfg.hostDriverVersion)
 	if err != nil {
 		return fmt.Errorf("failed to get container forward compat directory: %w", err)
 	}
@@ -119,17 +126,17 @@ func (m command) run(_ *cli.Command, cfg *options) error {
 	return m.createLdsoconfdFile(containerRoot(containerRootDir), cudaCompatLdsoconfdFilenamePattern, containerForwardCompatDir)
 }
 
-func (m command) getContainerForwardCompatDir(containerRoot containerRoot, hostDriverVersion string) (string, error) {
+func (m command) getContainerForwardCompatDir(containerRoot containerRoot, cudaCompatRoot string, hostDriverVersion string) (string, error) {
 	if hostDriverVersion == "" {
 		m.logger.Debugf("Host driver version not specified")
 		return "", nil
 	}
-	if !containerRoot.hasPath(cudaCompatPath) {
+	if !containerRoot.hasPath(cudaCompatRoot) {
 		m.logger.Debugf("No CUDA forward compatibility libraries directory in container")
 		return "", nil
 	}
 
-	libs, err := containerRoot.globFiles(filepath.Join(cudaCompatPath, "libcuda.so.*.*"))
+	libs, err := containerRoot.globFiles(filepath.Join(cudaCompatRoot, "libcuda.so.*.*"))
 	if err != nil {
 		m.logger.Warningf("Failed to find CUDA compat library: %w", err)
 		return "", nil

cmd/nvidia-cdi-hook/cudacompat/cudacompat_test.go

@@ -131,7 +131,7 @@ func TestCompatLibs(t *testing.T) {
 			c := command{
 				logger: logger,
 			}
-			containerForwardCompatDir, err := c.getContainerForwardCompatDir(containerRoot(containerRootDir), tc.hostDriverVersion)
+			containerForwardCompatDir, err := c.getContainerForwardCompatDir(containerRoot(containerRootDir), defaultCudaCompatPath, tc.hostDriverVersion)
 			require.NoError(t, err)
 			require.EqualValues(t, tc.expectedContainerForwardCompatDir, containerForwardCompatDir)
 		})

internal/discover/compat_libs.go

@@ -8,11 +8,14 @@ import (
 
 // 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, version string) Discover {
+func NewCUDACompatHookDiscoverer(logger logger.Interface, hookCreator HookCreator, version string, cudaCompatContainerRoot string) Discover {
 	var args []string
 	if version != "" && !strings.Contains(version, "*") {
 		args = append(args, "--host-driver-version="+version)
 	}
+	if cudaCompatContainerRoot != "" {
+		args = append(args, "--cuda-compat-container-root="+cudaCompatContainerRoot)
+	}
 
 	return hookCreator.Create("enable-cuda-compat", args...)
 }

internal/modifier/gated.go

@@ -107,7 +107,7 @@ func getCudaCompatModeDiscoverer(logger logger.Interface, cfg *config.Config, dr
 		return nil, fmt.Errorf("failed to get driver version: %w", err)
 	}
 
-	compatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(logger, hookCreator, version)
+	compatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(logger, hookCreator, version, "")
 	// For non-legacy modes we return the hook as is. These modes *should* already include the update-ldcache hook.
 	if cfg.NVIDIAContainerRuntimeConfig.Mode != "legacy" {
 		return compatLibHookDiscoverer, nil

internal/platform-support/tegra/tegra.go

@@ -80,6 +80,10 @@ func New(opts ...Option) (discover.Discover, error) {
 		return nil, fmt.Errorf("failed to create CSV discoverer: %v", err)
 	}
 
+	// TODO: Set the version here.
+	version := ""
+	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(o.logger, o.hookCreator, version)
+
 	ldcacheUpdateHook, err := discover.NewLDCacheUpdateHook(o.logger, csvDiscoverer, o.hookCreator, o.ldconfigPath)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create ldcach update hook discoverer: %v", err)
@@ -96,6 +100,7 @@ func New(opts ...Option) (discover.Discover, error) {
 
 	d := discover.Merge(
 		csvDiscoverer,
+		cudaCompatLibHookDiscoverer,
 		// The ldcacheUpdateHook is added last to ensure that the created symlinks are included
 		ldcacheUpdateHook,
 		tegraSystemMounts,

pkg/nvcdi/driver-nvml.go

@@ -101,7 +101,7 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string, libcudaSoParentDir
 	)
 	discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)
 
-	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, version)
+	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, version, "")
 	discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
 
 	updateLDCache, _ := discover.NewLDCacheUpdateHook(l.logger, libraries, l.hookCreator, l.ldconfigPath)