Add enable-cuda-compat on Tegra-based systems

This change ensures that tegra-based systems also include a
enable-cuda-compat hook. In the case of Orin devices, this
compat root in the container is set to /usr/local/cuda/compat-orin
to allow these compat libraries to be selected (if present in the
container).

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

Author: elezar

internal/platform-support/tegra/tegra.go

@@ -80,15 +80,6 @@ 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)
-	}
-
 	tegraSystemMounts := discover.NewMounts(
 		o.logger,
 		lookup.NewFileLocator(lookup.WithLogger(o.logger)),
@@ -100,9 +91,6 @@ 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/lib-csv.go

@@ -18,10 +18,14 @@ package nvcdi
 
 import (
 	"fmt"
+	"strings"
 
 	"tags.cncf.io/container-device-interface/pkg/cdi"
 	"tags.cncf.io/container-device-interface/specs-go"
 
+	"github.com/NVIDIA/go-nvlib/pkg/nvlib/device"
+	"github.com/NVIDIA/go-nvml/pkg/nvml"
+
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/discover"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/edits"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/platform-support/tegra"
@@ -46,7 +50,33 @@ func (l *csvlib) DeviceSpecGenerators(ids ...string) (DeviceSpecGenerator, error
 
 // GetDeviceSpecs returns the CDI device specs for a single device.
 func (l *csvlib) GetDeviceSpecs() ([]specs.Device, error) {
-	d, err := tegra.New(
+	d, err := l.driverDiscoverer()
+	if err != nil {
+		return nil, fmt.Errorf("failed to create driver discoverer from CSV files: %w", err)
+	}
+	e, err := edits.FromDiscoverer(d)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create container edits for CSV files: %w", err)
+	}
+
+	names, err := l.deviceNamers.GetDeviceNames(0, uuidIgnored{})
+	if err != nil {
+		return nil, fmt.Errorf("failed to get device name: %w", err)
+	}
+	var deviceSpecs []specs.Device
+	for _, name := range names {
+		deviceSpec := specs.Device{
+			Name:           name,
+			ContainerEdits: *e.ContainerEdits,
+		}
+		deviceSpecs = append(deviceSpecs, deviceSpec)
+	}
+
+	return deviceSpecs, nil
+}
+
+func (l *csvlib) driverDiscoverer() (discover.Discover, error) {
+	driverDiscoverer, err := tegra.New(
 		tegra.WithLogger(l.logger),
 		tegra.WithDriverRoot(l.driverRoot),
 		tegra.WithDevRoot(l.devRoot),
@@ -57,27 +87,76 @@ func (l *csvlib) GetDeviceSpecs() ([]specs.Device, error) {
 		tegra.WithIngorePatterns(l.csvIgnorePatterns...),
 	)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create discoverer for CSV files: %v", err)
+		return nil, fmt.Errorf("failed to create discoverer for CSV files: %w", err)
 	}
-	e, err := edits.FromDiscoverer(d)
+
+	cudaCompatDiscoverer := l.cudaCompatDiscoverer()
+
+	ldcacheUpdateHook, err := discover.NewLDCacheUpdateHook(l.logger, driverDiscoverer, l.hookCreator, l.ldconfigPath)
 	if err != nil {
-		return nil, fmt.Errorf("failed to create container edits for CSV files: %v", err)
+		return nil, fmt.Errorf("failed to create ldcache update hook discoverer: %w", err)
 	}
 
-	names, err := l.deviceNamers.GetDeviceNames(0, uuidIgnored{})
+	d := discover.Merge(
+		driverDiscoverer,
+		cudaCompatDiscoverer,
+		// The ldcacheUpdateHook is added last to ensure that the created symlinks are included
+		ldcacheUpdateHook,
+	)
+	return d, nil
+}
+
+// cudaCompatDiscoverer returns a discoverer for the CUDA forward compat hook
+// on Tegra-based systems.
+// It the system has NVML available, this is used to determine the driver
+// version to be passed to the hook.
+// On Orin-based systems, the compat library root in the container is also set.
+func (l *csvlib) cudaCompatDiscoverer() discover.Discover {
+	hasNvml, _ := l.infolib.HasNvml()
+	if !hasNvml {
+		return nil
+	}
+
+	ret := l.nvmllib.Init()
+	if ret != nvml.SUCCESS {
+		l.logger.Warningf("Failed to initialize NVML: %v", ret)
+		return nil
+	}
+	defer func() {
+		_ = l.nvmllib.Shutdown()
+	}()
+
+	version, ret := l.nvmllib.SystemGetDriverVersion()
+	if ret != nvml.SUCCESS {
+		l.logger.Warningf("Failed to get driver version: %v", ret)
+		return nil
+	}
+
+	var names []string
+	err := l.devicelib.VisitDevices(func(i int, d device.Device) error {
+		name, ret := d.GetName()
+		if ret != nvml.SUCCESS {
+			return fmt.Errorf("device %v: %v", i, ret)
+		}
+		names = append(names, name)
+		return nil
+	})
 	if err != nil {
-		return nil, fmt.Errorf("failed to get device name: %v", err)
+		l.logger.Warningf("Failed to get device names: %v", err)
+		return nil
 	}
-	var deviceSpecs []specs.Device
+
+	var cudaCompatContainerRoot string
 	for _, name := range names {
-		deviceSpec := specs.Device{
-			Name:           name,
-			ContainerEdits: *e.ContainerEdits,
+		// TODO: Should this be overridable through a feature flag / config option?
+		if strings.Contains(name, "Orin (nvgpu)") {
+			// TODO: This should probably be a constant or configurable.
+			cudaCompatContainerRoot = "/usr/local/cuda/compat-orin"
+			break
 		}
-		deviceSpecs = append(deviceSpecs, deviceSpec)
 	}
 
-	return deviceSpecs, nil
+	return discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, version, cudaCompatContainerRoot)
 }
 
 // GetCommonEdits generates a CDI specification that can be used for ANY devices