cmd, pkg/nvidia: Enable the proprietary NVIDIA driver

This uses the NVIDIA Container Toolkit [1] to generate a Container
Device Interface specification [2] on the host during the 'enter' and
'run' commands.  The specification is saved as JSON in the runtime
directories at /run/toolbox or $XDG_RUNTIME_DIR/toolbox to make it
available to the Toolbx container's entry point.  The environment
variables in the specification are directly passed to 'podman exec',
while the hooks and mounts are handled by the entry point.

Toolbx containers already have access to all the devices in the host
operating system's /dev, and containers share the kernel space driver
with the host.  So, this is only about making the user space driver
available to the container.  It's done by bind mounting the files
mentioned in the generated CDI specification from the host to the
container, and then updating the container's dynamic linker cache.

This neither depends on 'nvidia-ctk cdi generate' to generate the
Container Device Interface specification nor on 'podman create --device'
to consume it.

The main problem with nvidia-ctk and 'podman create' is that the
specification must be saved in /etc/cdi or /var/run/cdi, both of which
require root access, for it to be visible to 'podman create --device'.
Toolbx containers are often used rootless, so requiring root privileges
for hardware support, something that's not necessary on the host, will
be a problem.

Secondly, updating the toolbox(1) binary won't let existing containers
use the proprietary NVIDIA driver, because 'podman create' only affects
new containers.

Therefore, toolbox(1) uses the Go APIs used by 'nvidia-ctk cdi generate'
and 'podman create --device' to generate, save, load and apply the CDI
specification itself.  This removes the need for root privileges due to
/etc/cdi or /var/run/cdi, and makes the driver available to existing
containers.

Until Bats 1.10.0, 'run --keep-empty-lines' had a bug where it counted
the trailing newline on the last line as a separate line [3].  However,
Bats 1.10.0 is only available in Fedora >= 39 and is absent from Fedora
38.

Based on an idea from Ievgen Popovych.

[1] https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/
    https://github.com/NVIDIA/nvidia-container-toolkit

[2] https://github.com/cncf-tags/container-device-interface

[3] Bats commit 6648e2143bffb933
    bats-core/bats-core@6648e2143bffb933
    bats-core/bats-core#708

#116

Author: debarshiray

src/cmd/initContainer.go

@@ -30,9 +30,12 @@ import (
 	"github.com/containers/toolbox/pkg/shell"
 	"github.com/containers/toolbox/pkg/utils"
 	"github.com/fsnotify/fsnotify"
+	"github.com/google/renameio/v2"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.org/x/sys/unix"
+	"tags.cncf.io/container-device-interface/pkg/cdi"
+	"tags.cncf.io/container-device-interface/specs-go"
 )
 
 var (
@@ -264,6 +267,36 @@ func initContainer(cmd *cobra.Command, args []string) error {
 		return err
 	}
 
+	uidString := strconv.Itoa(initContainerFlags.uid)
+	targetUser, err := user.LookupId(uidString)
+	if err != nil {
+		return fmt.Errorf("failed to look up user ID %s: %w", uidString, err)
+	}
+
+	cdiFileForNvidia, err := getCDIFileForNvidia(targetUser)
+	if err != nil {
+		return err
+	}
+
+	logrus.Debugf("Loading Container Device Interface for NVIDIA from file %s", cdiFileForNvidia)
+
+	cdiSpecForNvidia, err := loadCDISpecFrom(cdiFileForNvidia)
+	if err != nil {
+		if errors.Is(err, os.ErrNotExist) {
+			logrus.Debugf("Loading Container Device Interface for NVIDIA: file %s not found",
+				cdiFileForNvidia)
+		} else {
+			logrus.Debugf("Loading Container Device Interface for NVIDIA: failed: %s", err)
+			return errors.New("failed to load Container Device Interface for NVIDIA")
+		}
+	}
+
+	if cdiSpecForNvidia != nil {
+		if err := applyCDISpecForNvidia(cdiSpecForNvidia); err != nil {
+			return err
+		}
+	}
+
 	if utils.PathExists("/etc/krb5.conf.d") && !utils.PathExists("/etc/krb5.conf.d/kcm_default_ccache") {
 		logrus.Debug("Setting KCM as the default Kerberos credential cache")
 
@@ -338,12 +371,6 @@ func initContainer(cmd *cobra.Command, args []string) error {
 
 	logrus.Debug("Finished initializing container")
 
-	uidString := strconv.Itoa(initContainerFlags.uid)
-	targetUser, err := user.LookupId(uidString)
-	if err != nil {
-		return fmt.Errorf("failed to look up user ID %s: %w", uidString, err)
-	}
-
 	toolboxRuntimeDirectory, err := utils.GetRuntimeDirectory(targetUser)
 	if err != nil {
 		return err
@@ -404,6 +431,83 @@ func initContainerHelp(cmd *cobra.Command, args []string) {
 	}
 }
 
+func applyCDISpecForNvidia(spec *specs.Spec) error {
+	if spec == nil {
+		panic("spec not specified")
+	}
+
+	logrus.Debug("Applying Container Device Interface for NVIDIA")
+
+	for _, mount := range spec.ContainerEdits.Mounts {
+		if err := (&cdi.Mount{Mount: mount}).Validate(); err != nil {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: invalid mount: %s", err)
+			return errors.New("invalid mount in Container Device Interface for NVIDIA")
+		}
+
+		if mount.Type == "" {
+			mount.Type = "bind"
+		}
+
+		if mount.Type != "bind" {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: unknown mount type %s",
+				mount.Type)
+			continue
+		}
+
+		flags := strings.Join(mount.Options, ",")
+		hostPath := filepath.Join(string(filepath.Separator), "run", "host", mount.HostPath)
+		if err := mountBind(mount.ContainerPath, hostPath, flags); err != nil {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: %s", err)
+			return errors.New("failed to apply mount from Container Device Interface for NVIDIA")
+		}
+	}
+
+	for _, hook := range spec.ContainerEdits.Hooks {
+		if err := (&cdi.Hook{Hook: hook}).Validate(); err != nil {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: invalid hook: %s", err)
+			return errors.New("invalid hook in Container Device Interface for NVIDIA")
+		}
+
+		if hook.HookName != cdi.CreateContainerHook {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: unknown hook name %s",
+				hook.HookName)
+			continue
+		}
+
+		if len(hook.Args) < 3 ||
+			hook.Args[0] != "nvidia-ctk" ||
+			hook.Args[1] != "hook" ||
+			hook.Args[2] != "update-ldcache" {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: unknown hook arguments")
+			continue
+		}
+
+		var folderFlag bool
+		var folders []string
+		hookArgs := hook.Args[3:]
+
+		for _, hookArg := range hookArgs {
+			if hookArg == "--folder" {
+				folderFlag = true
+				continue
+			}
+
+			if folderFlag {
+				folders = append(folders, hookArg)
+			}
+
+			folderFlag = false
+		}
+
+		if err := ldConfig("toolbx-nvidia.conf", folders); err != nil {
+			logrus.Debugf("Applying Container Device Interface for NVIDIA: %s", err)
+			return errors.New("failed to update ldcache for Container Device Interface for NVIDIA")
+		}
+	}
+
+	return nil
+}
+
 func configureUsers(targetUserUid int, targetUser, targetUserHome, targetUserShell string, homeLink bool) error {
 	if homeLink {
 		if err := redirectPath("/home", "/var/home", true); err != nil {
@@ -517,6 +621,73 @@ func handleFileSystemEvent(event fsnotify.Event) {
 	}
 }
 
+func ldConfig(configFileBase string, dirs []string) error {
+	logrus.Debug("Updating dynamic linker cache")
+
+	var args []string
+
+	if !utils.PathExists("/etc/ld.so.cache") {
+		logrus.Debug("Updating dynamic linker cache: no /etc/ld.so.cache found")
+		args = append(args, "-N")
+	}
+
+	if utils.PathExists("/etc/ld.so.conf.d") {
+		if len(dirs) > 0 {
+			var builder strings.Builder
+			builder.WriteString("# Written by Toolbx\n")
+			builder.WriteString("# https://containertoolbx.org/\n")
+			builder.WriteString("\n")
+
+			configured := make(map[string]struct{})
+
+			for _, dir := range dirs {
+				if _, ok := configured[dir]; ok {
+					continue
+				}
+
+				configured[dir] = struct{}{}
+				builder.WriteString(dir)
+				builder.WriteString("\n")
+			}
+
+			dirConfigString := builder.String()
+			dirConfigBytes := []byte(dirConfigString)
+			configFile := filepath.Join("/etc/ld.so.conf.d", configFileBase)
+			if err := renameio.WriteFile(configFile, dirConfigBytes, 0644); err != nil {
+				logrus.Debugf("Updating dynamic linker cache: failed to update configuration: %s", err)
+				return errors.New("failed to update dynamic linker cache configuration")
+			}
+		}
+	} else {
+		logrus.Debug("Updating dynamic linker cache: no /etc/ld.so.conf.d found")
+		args = append(args, dirs...)
+	}
+
+	if err := shell.Run("ldconfig", nil, nil, nil, args...); err != nil {
+		logrus.Debugf("Updating dynamic linker cache: failed: %s", err)
+		return errors.New("failed to update dynamic linker cache")
+	}
+
+	return nil
+}
+
+func loadCDISpecFrom(path string) (*specs.Spec, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+
+	spec, err := cdi.ParseSpec(data)
+	if err != nil {
+		return nil, err
+	}
+	if spec == nil {
+		return nil, errors.New("missing data")
+	}
+
+	return spec, nil
+}
+
 func mountBind(containerPath, source, flags string) error {
 	fi, err := os.Stat(source)
 	if err != nil {
@@ -537,6 +708,11 @@ func mountBind(containerPath, source, flags string) error {
 	} else if fileMode.IsRegular() {
 		logrus.Debugf("Creating regular file %s", containerPath)
 
+		containerPathDir := filepath.Dir(containerPath)
+		if err := os.MkdirAll(containerPathDir, 0755); err != nil {
+			return fmt.Errorf("failed to create directory %s: %w", containerPathDir, err)
+		}
+
 		containerPathFile, err := os.Create(containerPath)
 		if err != nil && !os.IsExist(err) {
 			return fmt.Errorf("failed to create regular file %s: %w", containerPath, err)

src/cmd/root.go

@@ -26,6 +26,7 @@ import (
 	"strings"
 	"syscall"
 
+	"github.com/containers/toolbox/pkg/nvidia"
 	"github.com/containers/toolbox/pkg/podman"
 	"github.com/containers/toolbox/pkg/utils"
 	"github.com/containers/toolbox/pkg/version"
@@ -382,6 +383,7 @@ func setUpLoggers() error {
 	logrus.SetLevel(logLevel)
 
 	if rootFlags.verbose > 1 {
+		nvidia.SetLogLevel(logLevel)
 		rootFlags.logPodman = true
 	}
 

src/cmd/run.go

@@ -19,6 +19,7 @@ package cmd
 import (
 	"bufio"
 	"context"
+	"encoding/json"
 	"errors"
 	"fmt"
 	"io"
@@ -28,15 +29,18 @@ import (
 	"strings"
 	"time"
 
+	"github.com/containers/toolbox/pkg/nvidia"
 	"github.com/containers/toolbox/pkg/podman"
 	"github.com/containers/toolbox/pkg/shell"
 	"github.com/containers/toolbox/pkg/term"
 	"github.com/containers/toolbox/pkg/utils"
 	"github.com/fsnotify/fsnotify"
 	"github.com/go-logfmt/logfmt"
+	"github.com/google/renameio/v2"
 	"github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
 	"golang.org/x/sys/unix"
+	"tags.cncf.io/container-device-interface/specs-go"
 )
 
 type collectEntryPointErrorFunc func(err error)
@@ -273,9 +277,31 @@ func runCommand(container string,
 		return err
 	}
 
+	var cdiEnviron []string
+
+	cdiSpecForNvidia, err := nvidia.GenerateCDISpec()
+	if err != nil {
+		if !errors.Is(err, nvidia.ErrPlatformUnsupported) {
+			return err
+		}
+	} else {
+		cdiEnviron = append(cdiEnviron, cdiSpecForNvidia.ContainerEdits.Env...)
+	}
+
 	startContainerTimestamp := time.Unix(-1, 0)
 
 	if entryPointPID <= 0 {
+		if cdiSpecForNvidia != nil {
+			cdiFileForNvidia, err := getCDIFileForNvidia(currentUser)
+			if err != nil {
+				return err
+			}
+
+			if err := saveCDISpecTo(cdiSpecForNvidia, cdiFileForNvidia); err != nil {
+				return err
+			}
+		}
+
 		startContainerTimestamp = time.Now()
 
 		logrus.Debugf("Starting container %s", container)
@@ -317,6 +343,7 @@ func runCommand(container string,
 	if err := runCommandWithFallbacks(container,
 		preserveFDs,
 		command,
+		cdiEnviron,
 		emitEscapeSequence,
 		fallbackToBash); err != nil {
 		return err
@@ -327,7 +354,7 @@ func runCommand(container string,
 
 func runCommandWithFallbacks(container string,
 	preserveFDs uint,
-	command []string,
+	command, environ []string,
 	emitEscapeSequence, fallbackToBash bool) error {
 
 	logrus.Debug("Checking if 'podman exec' supports disabling the detach keys")
@@ -340,6 +367,12 @@ func runCommandWithFallbacks(container string,
 	}
 
 	envOptions := utils.GetEnvOptionsForPreservedVariables()
+	for _, env := range environ {
+		logrus.Debugf("%s", env)
+		envOption := "--env=" + env
+		envOptions = append(envOptions, envOption)
+	}
+
 	preserveFDsString := fmt.Sprint(preserveFDs)
 
 	var stderr io.Writer
@@ -828,6 +861,36 @@ func isUsePollingSet() bool {
 	return true
 }
 
+func saveCDISpecTo(spec *specs.Spec, path string) error {
+	if path == "" {
+		panic("path not specified")
+	}
+
+	if spec == nil {
+		panic("spec not specified")
+	}
+
+	logrus.Debugf("Saving Container Device Interface to file %s", path)
+
+	if extension := filepath.Ext(path); extension != ".json" {
+		panicMsg := fmt.Sprintf("path has invalid extension %s", extension)
+		panic(panicMsg)
+	}
+
+	data, err := json.MarshalIndent(spec, "", "  ")
+	if err != nil {
+		logrus.Debugf("Saving Container Device Interface: failed to marshal JSON: %s", err)
+		return errors.New("failed to marshal Container Device Interface to JSON")
+	}
+
+	if err := renameio.WriteFile(path, data, 0644); err != nil {
+		logrus.Debugf("Saving Container Device Interface: failed to write file: %s", err)
+		return errors.New("failed to write Container Device Interface to file")
+	}
+
+	return nil
+}
+
 func showEntryPointLog(line string) error {
 	var logLevel logrus.Level
 	var logLevelFound bool