Add --cuda-compat-container-root

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

Author: elezar

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

@@ -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:       cudaCompatPath,
+				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,7 +126,7 @@ 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

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/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)