add default runtime binary path to runtimes field of toolkit config toml

Signed-off-by: Tariq Ibrahim <[email protected]>

Author: tariq1890

internal/config/toml.go

@@ -170,11 +170,26 @@ func (t *Toml) Get(key string) interface{} {
 	return (*toml.Tree)(t).Get(key)
 }
 
+// GetDefault returns the value for the specified key and falls back to the default value if the Get call fails
+func (t *Toml) GetDefault(key string, def interface{}) interface{} {
+	val := (*toml.Tree)(t).Get(key)
+	if val == nil {
+		return def
+	}
+	return val
+}
+
 // Set sets the specified key to the specified value in the TOML config.
 func (t *Toml) Set(key string, value interface{}) {
 	(*toml.Tree)(t).Set(key, value)
 }
 
+// WriteTo encode the Tree as Toml and writes it to the writer w.
+// Returns the number of bytes written in case of success, or an error if anything happened.
+func (t *Toml) WriteTo(w io.Writer) (int64, error) {
+	return (*toml.Tree)(t).WriteTo(w)
+}
+
 // commentDefaults applies the required comments for default values to the Toml.
 func (t *Toml) commentDefaults() *Toml {
 	asToml := (*toml.Tree)(t)

pkg/config/engine/api.go

@@ -23,4 +23,5 @@ type Interface interface {
 	Set(string, interface{})
 	RemoveRuntime(string) error
 	Save(string) (int64, error)
+	GetRuntime(name string) (interface{}, error)
 }

pkg/config/engine/containerd/config_v1.go

@@ -146,6 +146,18 @@ func (c *ConfigV1) RemoveRuntime(name string) error {
 	return nil
 }
 
+func (c *ConfigV1) GetRuntime(name string) (interface{}, error) {
+	if c == nil || c.Tree == nil {
+		return nil, fmt.Errorf("config is nil")
+	}
+	config := *c.Tree
+	runtimeData, ok := config.GetPath([]string{"plugins", "cri", "containerd", "runtimes", name}).(*toml.Tree)
+	if !ok {
+		return nil, fmt.Errorf("invalid toml object")
+	}
+	return runtimeData, nil
+}
+
 // SetOption sets the specified containerd option.
 func (c *ConfigV1) Set(key string, value interface{}) {
 	config := *c.Tree

pkg/config/engine/docker/docker.go

@@ -132,3 +132,25 @@ func (c Config) Save(path string) (int64, error) {
 	n, err := config.Raw(path).Write(output)
 	return int64(n), err
 }
+
+// GetRuntime returns the TOML object of the runtime passed as input
+func (c *Config) GetRuntime(name string) (interface{}, error) {
+	if c == nil {
+		return nil, fmt.Errorf("config is nil")
+	}
+
+	config := *c
+
+	var runtimes map[string]interface{}
+	if _, exists := config["runtimes"]; exists {
+		runtimes = config["runtimes"].(map[string]interface{})
+	} else {
+		return nil, fmt.Errorf("config is nil")
+	}
+
+	if r, ok := runtimes[name]; ok {
+		return r, nil
+	} else {
+		return nil, fmt.Errorf("runtime %s not found", name)
+	}
+}

tools/container/toolkit/toolkit.go

@@ -32,6 +32,9 @@ import (
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/config"
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/system/nvdevices"
+	pkgconfig "github.com/NVIDIA/nvidia-container-toolkit/pkg/config"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/containerd"
+	"github.com/NVIDIA/nvidia-container-toolkit/pkg/config/engine/crio"
 	"github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi"
 	transformroot "github.com/NVIDIA/nvidia-container-toolkit/pkg/nvcdi/transform/root"
 )
@@ -40,6 +43,8 @@ const (
 	// DefaultNvidiaDriverRoot specifies the default NVIDIA driver run directory
 	DefaultNvidiaDriverRoot = "/run/nvidia/driver"
 
+	DefaultHostRootMount = "/host"
+
 	nvidiaContainerCliSource         = "/usr/bin/nvidia-container-cli"
 	nvidiaContainerRuntimeHookSource = "/usr/bin/nvidia-container-runtime-hook"
 
@@ -54,6 +59,7 @@ type Options struct {
 	DevRoot           string
 	DriverRootCtrPath string
 	DevRootCtrPath    string
+	HostRootMount     string
 
 	ContainerRuntimeMode     string
 	ContainerRuntimeDebug    string
@@ -109,6 +115,13 @@ func Flags(opts *Options) []cli.Flag {
 			Destination: &opts.DevRootCtrPath,
 			EnvVars:     []string{"DEV_ROOT_CTR_PATH"},
 		},
+		&cli.StringFlag{
+			Name:        "host-root",
+			Usage:       "Specify the path to the host root to be used when executing shell commands.",
+			Value:       DefaultHostRootMount,
+			Destination: &opts.HostRootMount,
+			EnvVars:     []string{"HOST_ROOT_MOUNT"},
+		},
 		&cli.StringFlag{
 			Name:        "nvidia-container-runtime.debug",
 			Aliases:     []string{"nvidia-container-runtime-debug"},
@@ -342,7 +355,14 @@ func Install(cli *cli.Context, opts *Options, toolkitRoot string) error {
 		log.Errorf("Ignoring error: %v", fmt.Errorf("error installing NVIDIA Container CDI Hook CLI: %v", err))
 	}
 
-	err = installToolkitConfig(cli, toolkitConfigPath, nvidiaContainerCliExecutable, nvidiaCTKPath, nvidiaContainerRuntimeHookPath, opts)
+	var runtimeBinPaths []string
+	runtimeBinPaths, err = getRuntimeBinaryPaths(opts.HostRootMount)
+	if err != nil {
+		log.Warningf("Error retrieving runtime binary paths: %v", err)
+	}
+
+	err = installToolkitConfig(cli, toolkitConfigPath, nvidiaContainerCliExecutable, nvidiaCTKPath,
+		nvidiaContainerRuntimeHookPath, runtimeBinPaths, opts)
 	if err != nil && !opts.ignoreErrors {
 		return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err)
 	} else if err != nil {
@@ -419,11 +439,15 @@ func installLibrary(libName string, toolkitRoot string) error {
 func installToolkitConfig(c *cli.Context, toolkitConfigPath string, nvidiaContainerCliExecutablePath string, nvidiaCTKPath string, nvidaContainerRuntimeHookPath string, opts *Options) error {
 	log.Infof("Installing NVIDIA container toolkit config '%v'", toolkitConfigPath)
 
-	cfg, err := loadConfig(nvidiaContainerToolkitConfigSource)
+	cfg, err := config.New(
+		config.WithConfigFile(opts.toolkitConfigSource),
+	)
 	if err != nil {
 		return fmt.Errorf("could not open source config file: %v", err)
 	}
 
+	log.Infof("Printing Default toolkit toml config: %v", cfg)
+
 	targetConfig, err := os.Create(toolkitConfigPath)
 	if err != nil {
 		return fmt.Errorf("could not create target config file: %v", err)
@@ -436,6 +460,18 @@ func installToolkitConfig(c *cli.Context, toolkitConfigPath string, nvidiaContai
 	// Use the driver run root as the root:
 	driverLdconfigPath := config.NormalizeLDConfigPath("@" + filepath.Join(opts.DriverRoot, strings.TrimPrefix(ldconfigPath, "@/")))
 
+	var ctkRuntimes []string
+	defaultCfg, err := cfg.Config()
+	if err == nil {
+		defaultCfgRuntimes := defaultCfg.NVIDIAContainerRuntimeConfig.Runtimes
+		if len(runtimeBinaryPaths) > 0 {
+			ctkRuntimes = append(ctkRuntimes, runtimeBinaryPaths...)
+		}
+		ctkRuntimes = append(ctkRuntimes, defaultCfgRuntimes...)
+	} else {
+		log.Warningf("could not get default toolkit config: %v", err)
+	}
+
 	configValues := map[string]interface{}{
 		// Set the options in the root toml table
 		"accept-nvidia-visible-devices-envvar-when-unprivileged": opts.acceptNVIDIAVisibleDevicesWhenUnprivileged,
@@ -450,6 +486,11 @@ func installToolkitConfig(c *cli.Context, toolkitConfigPath string, nvidiaContai
 		"nvidia-container-runtime-hook.path":                nvidaContainerRuntimeHookPath,
 		"nvidia-container-runtime-hook.skip-mode-detection": opts.ContainerRuntimeHookSkipModeDetection,
 	}
+
+	if len(ctkRuntimes) > 0 {
+		configValues["nvidia-container-runtime.runtimes"] = ctkRuntimes
+	}
+
 	for key, value := range configValues {
 		cfg.Set(key, value)
 	}
@@ -503,16 +544,6 @@ func installToolkitConfig(c *cli.Context, toolkitConfigPath string, nvidiaContai
 	return nil
 }
 
-func loadConfig(path string) (*toml.Tree, error) {
-	_, err := os.Stat(path)
-	if err == nil {
-		return toml.LoadFile(path)
-	} else if os.IsNotExist(err) {
-		return toml.TreeFromMap(nil)
-	}
-	return nil, err
-}
-
 // installContainerToolkitCLI installs the nvidia-ctk CLI executable and wrapper.
 func installContainerToolkitCLI(toolkitDir string) (string, error) {
 	e := executable{
@@ -793,3 +824,64 @@ func generateCDISpec(opts *Options, nvidiaCDIHookPath string) error {
 
 	return nil
 }
+
+// getRuntimeBinaryPaths extracts the full paths of the low-level runtime binaries specified in the container runtime config
+func getRuntimeBinaryPaths(hostRoot string) ([]string, error) {
+	var runtimeBinaryPaths []string
+
+	if r, ok := os.LookupEnv("RUNTIME"); ok {
+
+		switch r {
+		case "containerd":
+			var cliArgs []string
+			if hostRoot != "" {
+				cliArgs = append(cliArgs, "chroot", hostRoot)
+			}
+			cliArgs = append(cliArgs, "containerd", "config", "dump")
+
+			cfg, err := containerd.New(
+				containerd.WithReference(pkgconfig.FromCLI(cliArgs)),
+			)
+			if err != nil {
+				return nil, fmt.Errorf("unable to load containerd config: %w", err)
+			}
+
+			containerdDefaultRuntime := cfg.DefaultRuntime()
+			if containerdDefaultRuntime != "" {
+				runtimeData, err := cfg.GetRuntime(containerdDefaultRuntime)
+				if err == nil {
+					runtimeTOML := runtimeData.(*toml.Tree)
+					if binaryPath, ok := runtimeTOML.GetPath([]string{"options", "BinaryName"}).(string); ok && binaryPath != "" {
+						runtimeBinaryPaths = append(runtimeBinaryPaths, binaryPath)
+					}
+				}
+			}
+
+		case "crio":
+			var cliArgs []string
+			if hostRoot != "" {
+				cliArgs = append(cliArgs, "chroot", hostRoot)
+			}
+			cliArgs = append(cliArgs, "crio", "status", "config")
+
+			cfg, err := crio.New(
+				crio.WithReference(pkgconfig.FromCLI(cliArgs)),
+			)
+			if err != nil {
+				return nil, fmt.Errorf("unable to load crio config: %w", err)
+			}
+
+			crioDefaultRuntime := cfg.DefaultRuntime()
+			if crioDefaultRuntime != "" {
+				runtimeData, err := cfg.GetRuntime(crioDefaultRuntime)
+				if err == nil {
+					runtimeTOML := runtimeData.(*toml.Tree)
+					if binaryPath, ok := runtimeTOML.GetPath([]string{"runtime_path"}).(string); ok && binaryPath != "" {
+						runtimeBinaryPaths = append(runtimeBinaryPaths, binaryPath)
+					}
+				}
+			}
+		}
+	}
+	return runtimeBinaryPaths, nil
+}