Copy cmd from container-config

This change copies the code from container-config/cmd to tools/container. This allows the code to be built and added to the container image without additional refactoring. As the configuration utilities are incorporated into the cmds of the nvidia-container-toolkit, the code will be moved from tools. Files copied from: 383587f766 Signed-off-by: Evan Lezar <elezar@nvidia.com>
2025-06-26 18:18:24 +00:00 · 2021-10-11 16:31:02 +02:00
parent 7bf5c25831
commit 7d76243783
18 changed files with 4144 additions and 0 deletions
--- a/tools/container/toolkit/executable.go
+++ b/tools/container/toolkit/executable.go
@@ -0,0 +1,153 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	log "github.com/sirupsen/logrus"
+)
+
+type executableTarget struct {
+	dotfileName string
+	wrapperName string
+}
+
+type executable struct {
+	source   string
+	target   executableTarget
+	env      map[string]string
+	preLines []string
+	argLines []string
+}
+
+// install installs an executable component of the NVIDIA container toolkit. The source executable
+// is copied to a `.real` file and a wapper is created to set up the environment as required.
+func (e executable) install(destFolder string) (string, error) {
+	log.Infof("Installing executable '%v' to %v", e.source, destFolder)
+
+	dotfileName := e.dotfileName()
+
+	installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source)
+	if err != nil {
+		return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotfileName, err)
+	}
+	log.Infof("Installed '%v'", installedDotfileName)
+
+	wrapperFilename, err := e.installWrapper(destFolder, installedDotfileName)
+	if err != nil {
+		return "", fmt.Errorf("error wrapping '%v': %v", installedDotfileName, err)
+	}
+	log.Infof("Installed wrapper '%v'", wrapperFilename)
+
+	return wrapperFilename, nil
+}
+
+func (e executable) dotfileName() string {
+	return e.target.dotfileName
+}
+
+func (e executable) wrapperName() string {
+	return e.target.wrapperName
+}
+
+func (e executable) installWrapper(destFolder string, dotfileName string) (string, error) {
+	wrapperPath := filepath.Join(destFolder, e.wrapperName())
+	wrapper, err := os.Create(wrapperPath)
+	if err != nil {
+		return "", fmt.Errorf("error creating executable wrapper: %v", err)
+	}
+	defer wrapper.Close()
+
+	err = e.writeWrapperTo(wrapper, destFolder, dotfileName)
+	if err != nil {
+		return "", fmt.Errorf("error writing wrapper contents: %v", err)
+	}
+
+	err = ensureExecutable(wrapperPath)
+	if err != nil {
+		return "", fmt.Errorf("error making wrapper executable: %v", err)
+	}
+	return wrapperPath, nil
+}
+
+func (e executable) writeWrapperTo(wrapper io.Writer, destFolder string, dotfileName string) error {
+	r := newReplacements(destDirPattern, destFolder)
+
+	// Add the shebang
+	fmt.Fprintln(wrapper, "#! /bin/sh")
+
+	// Add the preceding lines if any
+	for _, line := range e.preLines {
+		fmt.Fprintf(wrapper, "%s\n", r.apply(line))
+	}
+
+	// Update the path to include the destination folder
+	var env map[string]string
+	if e.env == nil {
+		env = make(map[string]string)
+	} else {
+		env = e.env
+	}
+
+	path, specified := env["PATH"]
+	if !specified {
+		path = "$PATH"
+	}
+	env["PATH"] = strings.Join([]string{destFolder, path}, ":")
+
+	var sortedEnvvars []string
+	for e := range env {
+		sortedEnvvars = append(sortedEnvvars, e)
+	}
+	sort.Strings(sortedEnvvars)
+
+	for _, e := range sortedEnvvars {
+		v := env[e]
+		fmt.Fprintf(wrapper, "%s=%s \\\n", e, r.apply(v))
+	}
+	// Add the call to the target executable
+	fmt.Fprintf(wrapper, "%s \\\n", dotfileName)
+
+	// Insert additional lines in the `arg` list
+	for _, line := range e.argLines {
+		fmt.Fprintf(wrapper, "\t%s \\\n", r.apply(line))
+	}
+	// Add the script arguments "$@"
+	fmt.Fprintln(wrapper, "\t\"$@\"")
+
+	return nil
+}
+
+// ensureExecutable is equivalent to running chmod +x on the specified file
+func ensureExecutable(path string) error {
+	info, err := os.Stat(path)
+	if err != nil {
+		return fmt.Errorf("error getting file info for '%v': %v", path, err)
+	}
+	executableMode := info.Mode() | 0111
+	err = os.Chmod(path, executableMode)
+	if err != nil {
+		return fmt.Errorf("error setting executable mode for '%v': %v", path, err)
+	}
+	return nil
+}
--- a/tools/container/toolkit/executable_test.go
+++ b/tools/container/toolkit/executable_test.go
@@ -0,0 +1,152 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import (
+	"bytes"
+	"os"
+	"path/filepath"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWrapper(t *testing.T) {
+	const shebang = "#! /bin/sh"
+	const destFolder = "/dest/folder"
+	const dotfileName = "source.real"
+
+	testCases := []struct {
+		e             executable
+		expectedLines []string
+	}{
+		{
+			e: executable{},
+			expectedLines: []string{
+				shebang,
+				"PATH=/dest/folder:$PATH \\",
+				"source.real \\",
+				"\t\"$@\"",
+				"",
+			},
+		},
+		{
+			e: executable{
+				env: map[string]string{
+					"PATH": "some-path",
+				},
+			},
+			expectedLines: []string{
+				shebang,
+				"PATH=/dest/folder:some-path \\",
+				"source.real \\",
+				"\t\"$@\"",
+				"",
+			},
+		},
+		{
+			e: executable{
+				preLines: []string{
+					"preline1",
+					"preline2",
+				},
+			},
+			expectedLines: []string{
+				shebang,
+				"preline1",
+				"preline2",
+				"PATH=/dest/folder:$PATH \\",
+				"source.real \\",
+				"\t\"$@\"",
+				"",
+			},
+		},
+		{
+			e: executable{
+				argLines: []string{
+					"argline1",
+					"argline2",
+				},
+			},
+			expectedLines: []string{
+				shebang,
+				"PATH=/dest/folder:$PATH \\",
+				"source.real \\",
+				"\targline1 \\",
+				"\targline2 \\",
+				"\t\"$@\"",
+				"",
+			},
+		},
+	}
+
+	for i, tc := range testCases {
+		buf := &bytes.Buffer{}
+
+		err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
+		require.NoError(t, err)
+
+		exepectedContents := strings.Join(tc.expectedLines, "\n")
+		require.Equal(t, exepectedContents, buf.String(), "%v: %v", i, tc)
+	}
+}
+
+func TestInstallExecutable(t *testing.T) {
+	inputFolder, err := os.MkdirTemp("", "")
+	require.NoError(t, err)
+	defer os.RemoveAll(inputFolder)
+
+	// Create the source file
+	source := filepath.Join(inputFolder, "input")
+	sourceFile, err := os.Create(source)
+
+	base := filepath.Base(source)
+
+	require.NoError(t, err)
+	require.NoError(t, sourceFile.Close())
+
+	e := executable{
+		source: source,
+		target: executableTarget{
+			dotfileName: "input.real",
+			wrapperName: "input",
+		},
+	}
+
+	destFolder, err := os.MkdirTemp("", "output-*")
+	require.NoError(t, err)
+	defer os.RemoveAll(destFolder)
+
+	installed, err := e.install(destFolder)
+
+	require.NoError(t, err)
+	require.Equal(t, filepath.Join(destFolder, base), installed)
+
+	// Now check the post conditions:
+	sourceInfo, err := os.Stat(source)
+	require.NoError(t, err)
+
+	destInfo, err := os.Stat(filepath.Join(destFolder, base+".real"))
+	require.NoError(t, err)
+	require.Equal(t, sourceInfo.Size(), destInfo.Size())
+	require.Equal(t, sourceInfo.Mode(), destInfo.Mode())
+
+	wrapperInfo, err := os.Stat(installed)
+	require.NoError(t, err)
+	require.NotEqual(t, 0, wrapperInfo.Mode()&0111)
+}
--- a/tools/container/toolkit/replacements.go
+++ b/tools/container/toolkit/replacements.go
@@ -0,0 +1,45 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import "strings"
+
+const (
+	destDirPattern = "@destDir@"
+)
+
+type replacements map[string]string
+
+func newReplacements(rules ...string) replacements {
+	r := make(replacements)
+	for i := 0; i < len(rules)-1; i += 2 {
+		old := rules[i]
+		new := rules[i+1]
+
+		r[old] = new
+	}
+
+	return r
+}
+
+func (r replacements) apply(input string) string {
+	output := input
+	for old, new := range r {
+		output = strings.ReplaceAll(output, old, new)
+	}
+	return output
+}
--- a/tools/container/toolkit/runtime.go
+++ b/tools/container/toolkit/runtime.go
@@ -0,0 +1,132 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import (
+	"fmt"
+	"path/filepath"
+	"strings"
+
+	log "github.com/sirupsen/logrus"
+)
+
+const (
+	nvidiaContainerRuntimeSource  = "/usr/bin/nvidia-container-runtime"
+	nvidiaContainerRuntimeTarget  = "nvidia-container-runtime.real"
+	nvidiaContainerRuntimeWrapper = "nvidia-container-runtime"
+
+	nvidiaExperimentalContainerRuntimeSource  = "nvidia-container-runtime.experimental"
+	nvidiaExperimentalContainerRuntimeTarget  = nvidiaExperimentalContainerRuntimeSource
+	nvidiaExperimentalContainerRuntimeWrapper = "nvidia-container-runtime-experimental"
+)
+
+// installContainerRuntimes sets up the NVIDIA container runtimes, copying the executables
+// and implementing the required wrapper
+func installContainerRuntimes(toolkitDir string, driverRoot string) error {
+	r := newNvidiaContainerRuntimeInstaller()
+
+	_, err := r.install(toolkitDir)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container runtime: %v", err)
+	}
+
+	// Install the experimental runtime and treat failures as non-fatal.
+	err = installExperimentalRuntime(toolkitDir, driverRoot)
+	if err != nil {
+		log.Warnf("Could not install experimental runtime: %v", err)
+	}
+
+	return nil
+}
+
+// installExperimentalRuntime ensures that the experimental NVIDIA Container runtime is installed
+func installExperimentalRuntime(toolkitDir string, driverRoot string) error {
+	libraryRoot, err := findLibraryRoot(driverRoot)
+	if err != nil {
+		log.Warnf("Error finding library path for root %v: %v", driverRoot, err)
+	}
+	log.Infof("Using library root %v", libraryRoot)
+
+	e := newNvidiaContainerRuntimeExperimentalInstaller(libraryRoot)
+	_, err = e.install(toolkitDir)
+	if err != nil {
+		return fmt.Errorf("error installing experimental NVIDIA Container Runtime: %v", err)
+	}
+
+	return nil
+}
+
+func newNvidiaContainerRuntimeInstaller() *executable {
+	target := executableTarget{
+		dotfileName: nvidiaContainerRuntimeTarget,
+		wrapperName: nvidiaContainerRuntimeWrapper,
+	}
+	return newRuntimeInstaller(nvidiaContainerRuntimeSource, target, nil)
+}
+
+func newNvidiaContainerRuntimeExperimentalInstaller(libraryRoot string) *executable {
+	target := executableTarget{
+		dotfileName: nvidiaExperimentalContainerRuntimeTarget,
+		wrapperName: nvidiaExperimentalContainerRuntimeWrapper,
+	}
+
+	env := make(map[string]string)
+	if libraryRoot != "" {
+		env["LD_LIBRARY_PATH"] = strings.Join([]string{libraryRoot, "$LD_LIBRARY_PATH"}, ":")
+	}
+	return newRuntimeInstaller(nvidiaExperimentalContainerRuntimeSource, target, env)
+}
+
+func newRuntimeInstaller(source string, target executableTarget, env map[string]string) *executable {
+	preLines := []string{
+		"",
+		"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
+		"if [ \"${?}\" != \"0\" ]; then",
+		"	echo \"nvidia driver modules are not yet loaded, invoking runc directly\"",
+		"	exec runc \"$@\"",
+		"fi",
+		"",
+	}
+
+	runtimeEnv := make(map[string]string)
+	runtimeEnv["XDG_CONFIG_HOME"] = filepath.Join(destDirPattern, ".config")
+	for k, v := range env {
+		runtimeEnv[k] = v
+	}
+
+	r := executable{
+		source:   source,
+		target:   target,
+		env:      runtimeEnv,
+		preLines: preLines,
+	}
+
+	return &r
+}
+
+func findLibraryRoot(root string) (string, error) {
+	libnvidiamlPath, err := findManagementLibrary(root)
+	if err != nil {
+		return "", fmt.Errorf("error locating NVIDIA management library: %v", err)
+	}
+
+	return filepath.Dir(libnvidiamlPath), nil
+}
+
+func findManagementLibrary(root string) (string, error) {
+	return findLibrary(root, "libnvidia-ml.so")
+}
--- a/tools/container/toolkit/runtime_test.go
+++ b/tools/container/toolkit/runtime_test.go
@@ -0,0 +1,90 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import (
+	"bytes"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestNvidiaContainerRuntimeInstallerWrapper(t *testing.T) {
+	r := newNvidiaContainerRuntimeInstaller()
+
+	const shebang = "#! /bin/sh"
+	const destFolder = "/dest/folder"
+	const dotfileName = "source.real"
+
+	buf := &bytes.Buffer{}
+
+	err := r.writeWrapperTo(buf, destFolder, dotfileName)
+	require.NoError(t, err)
+
+	expectedLines := []string{
+		shebang,
+		"",
+		"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
+		"if [ \"${?}\" != \"0\" ]; then",
+		"	echo \"nvidia driver modules are not yet loaded, invoking runc directly\"",
+		"	exec runc \"$@\"",
+		"fi",
+		"",
+		"PATH=/dest/folder:$PATH \\",
+		"XDG_CONFIG_HOME=/dest/folder/.config \\",
+		"source.real \\",
+		"\t\"$@\"",
+		"",
+	}
+
+	exepectedContents := strings.Join(expectedLines, "\n")
+	require.Equal(t, exepectedContents, buf.String())
+}
+
+func TestExperimentalContainerRuntimeInstallerWrapper(t *testing.T) {
+	r := newNvidiaContainerRuntimeExperimentalInstaller("/some/root/usr/lib64")
+
+	const shebang = "#! /bin/sh"
+	const destFolder = "/dest/folder"
+	const dotfileName = "source.real"
+
+	buf := &bytes.Buffer{}
+
+	err := r.writeWrapperTo(buf, destFolder, dotfileName)
+	require.NoError(t, err)
+
+	expectedLines := []string{
+		shebang,
+		"",
+		"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
+		"if [ \"${?}\" != \"0\" ]; then",
+		"	echo \"nvidia driver modules are not yet loaded, invoking runc directly\"",
+		"	exec runc \"$@\"",
+		"fi",
+		"",
+		"LD_LIBRARY_PATH=/some/root/usr/lib64:$LD_LIBRARY_PATH \\",
+		"PATH=/dest/folder:$PATH \\",
+		"XDG_CONFIG_HOME=/dest/folder/.config \\",
+		"source.real \\",
+		"\t\"$@\"",
+		"",
+	}
+
+	exepectedContents := strings.Join(expectedLines, "\n")
+	require.Equal(t, exepectedContents, buf.String())
+}
--- a/tools/container/toolkit/toolkit.go
+++ b/tools/container/toolkit/toolkit.go
@@ -0,0 +1,449 @@
+/**
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+*/
+
+package main
+
+import (
+	"fmt"
+	"io"
+	"os"
+	"path/filepath"
+	"strings"
+
+	toml "github.com/pelletier/go-toml"
+	log "github.com/sirupsen/logrus"
+	"github.com/urfave/cli/v2"
+)
+
+const (
+	// DefaultNvidiaDriverRoot specifies the default NVIDIA driver run directory
+	DefaultNvidiaDriverRoot = "/run/nvidia/driver"
+
+	nvidiaContainerCliSource         = "/usr/bin/nvidia-container-cli"
+	nvidiaContainerRuntimeHookSource = "/usr/bin/nvidia-container-toolkit"
+
+	nvidiaContainerToolkitConfigSource = "/etc/nvidia-container-runtime/config.toml"
+	configFilename                     = "config.toml"
+)
+
+var toolkitDirArg string
+var nvidiaDriverRootFlag string
+var nvidiaContainerRuntimeDebugFlag string
+var nvidiaContainerRuntimeLogLevelFlag string
+var nvidiaContainerCLIDebugFlag string
+
+func main() {
+	// Create the top-level CLI
+	c := cli.NewApp()
+	c.Name = "toolkit"
+	c.Usage = "Manage the NVIDIA container toolkit"
+	c.Version = "0.1.0"
+
+	// Create the 'install' subcommand
+	install := cli.Command{}
+	install.Name = "install"
+	install.Usage = "Install the components of the NVIDIA container toolkit"
+	install.ArgsUsage = "<toolkit_directory>"
+	install.Before = parseArgs
+	install.Action = Install
+
+	// Create the 'delete' command
+	delete := cli.Command{}
+	delete.Name = "delete"
+	delete.Usage = "Delete the NVIDIA container toolkit"
+	delete.ArgsUsage = "<toolkit_directory>"
+	delete.Before = parseArgs
+	delete.Action = Delete
+
+	// Register the subcommand with the top-level CLI
+	c.Commands = []*cli.Command{
+		&install,
+		&delete,
+	}
+
+	flags := []cli.Flag{
+		&cli.StringFlag{
+			Name:        "nvidia-driver-root",
+			Value:       DefaultNvidiaDriverRoot,
+			Destination: &nvidiaDriverRootFlag,
+			EnvVars:     []string{"NVIDIA_DRIVER_ROOT"},
+		},
+		&cli.StringFlag{
+			Name:        "nvidia-container-runtime-debug",
+			Usage:       "Specify the location of the debug log file for the NVIDIA Container Runtime",
+			Destination: &nvidiaContainerRuntimeDebugFlag,
+			EnvVars:     []string{"NVIDIA_CONTAINER_RUNTIME_DEBUG"},
+		},
+		&cli.StringFlag{
+			Name:        "nvidia-container-runtime-debug-log-level",
+			Destination: &nvidiaContainerRuntimeLogLevelFlag,
+			EnvVars:     []string{"NVIDIA_CONTAINER_RUNTIME_LOG_LEVEL"},
+		},
+		&cli.StringFlag{
+			Name:        "nvidia-container-cli-debug",
+			Usage:       "Specify the location of the debug log file for the NVIDIA Container CLI",
+			Destination: &nvidiaContainerCLIDebugFlag,
+			EnvVars:     []string{"NVIDIA_CONTAINER_CLI_DEBUG"},
+		},
+	}
+
+	// Update the subcommand flags with the common subcommand flags
+	install.Flags = append([]cli.Flag{}, flags...)
+
+	// Run the top-level CLI
+	if err := c.Run(os.Args); err != nil {
+		log.Fatal(fmt.Errorf("error: %v", err))
+	}
+}
+
+// parseArgs parses the command line arguments to the CLI
+func parseArgs(c *cli.Context) error {
+	args := c.Args()
+
+	log.Infof("Parsing arguments: %v", args.Slice())
+	if c.NArg() != 1 {
+		return fmt.Errorf("incorrect number of arguments")
+	}
+	toolkitDirArg = args.Get(0)
+	log.Infof("Successfully parsed arguments")
+
+	return nil
+}
+
+// Delete removes the NVIDIA container toolkit
+func Delete(cli *cli.Context) error {
+	log.Infof("Deleting NVIDIA container toolkit from '%v'", toolkitDirArg)
+	err := os.RemoveAll(toolkitDirArg)
+	if err != nil {
+		return fmt.Errorf("error deleting toolkit directory: %v", err)
+	}
+	return nil
+}
+
+// Install installs the components of the NVIDIA container toolkit.
+// Any existing installation is removed.
+func Install(cli *cli.Context) error {
+	log.Infof("Installing NVIDIA container toolkit to '%v'", toolkitDirArg)
+
+	log.Infof("Removing existing NVIDIA container toolkit installation")
+	err := os.RemoveAll(toolkitDirArg)
+	if err != nil {
+		return fmt.Errorf("error removing toolkit directory: %v", err)
+	}
+
+	toolkitConfigDir := filepath.Join(toolkitDirArg, ".config", "nvidia-container-runtime")
+	toolkitConfigPath := filepath.Join(toolkitConfigDir, configFilename)
+
+	err = createDirectories(toolkitDirArg, toolkitConfigDir)
+	if err != nil {
+		return fmt.Errorf("could not create required directories: %v", err)
+	}
+
+	err = installContainerLibrary(toolkitDirArg)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container library: %v", err)
+	}
+
+	err = installContainerRuntimes(toolkitDirArg, nvidiaDriverRootFlag)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container runtime: %v", err)
+	}
+
+	nvidiaContainerCliExecutable, err := installContainerCLI(toolkitDirArg)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container CLI: %v", err)
+	}
+
+	_, err = installRuntimeHook(toolkitDirArg, toolkitConfigPath)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container runtime hook: %v", err)
+	}
+
+	err = installToolkitConfig(toolkitConfigPath, nvidiaDriverRootFlag, nvidiaContainerCliExecutable)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err)
+	}
+
+	return nil
+}
+
+// installContainerLibrary locates and installs the libnvidia-container.so.1 library.
+// A predefined set of library candidates are considered, with the first one
+// resulting in success being installed to the toolkit folder. The install process
+// resolves the symlink for the library and copies the versioned library itself.
+func installContainerLibrary(toolkitDir string) error {
+	log.Infof("Installing NVIDIA container library to '%v'", toolkitDir)
+
+	const libName = "libnvidia-container.so.1"
+	libraryPath, err := findLibrary("", libName)
+	if err != nil {
+		return fmt.Errorf("error locating NVIDIA container library: %v", err)
+	}
+
+	installedLibPath, err := installFileToFolder(toolkitDir, libraryPath)
+	if err != nil {
+		return fmt.Errorf("error installing %v to %v: %v", libraryPath, toolkitDir, err)
+	}
+	log.Infof("Installed '%v' to '%v'", libraryPath, installedLibPath)
+
+	if filepath.Base(installedLibPath) == libName {
+		return nil
+	}
+
+	err = installSymlink(toolkitDir, libName, installedLibPath)
+	if err != nil {
+		return fmt.Errorf("error installing symlink for NVIDIA container library: %v", err)
+	}
+
+	return nil
+}
+
+// installToolkitConfig installs the config file for the NVIDIA container toolkit ensuring
+// that the settings are updated to match the desired install and nvidia driver directories.
+func installToolkitConfig(toolkitConfigPath string, nvidiaDriverDir string, nvidiaContainerCliExecutablePath string) error {
+	log.Infof("Installing NVIDIA container toolkit config '%v'", toolkitConfigPath)
+
+	config, err := toml.LoadFile(nvidiaContainerToolkitConfigSource)
+	if err != nil {
+		return fmt.Errorf("could not open source config file: %v", err)
+	}
+
+	targetConfig, err := os.Create(toolkitConfigPath)
+	if err != nil {
+		return fmt.Errorf("could not create target config file: %v", err)
+	}
+	defer targetConfig.Close()
+
+	nvidiaContainerCliKey := func(p string) []string {
+		return []string{"nvidia-container-cli", p}
+	}
+
+	// Read the ldconfig path from the config as this may differ per platform
+	// On ubuntu-based systems this ends in `.real`
+	ldconfigPath := fmt.Sprintf("%s", config.GetPath(nvidiaContainerCliKey("ldconfig")))
+
+	// Use the driver run root as the root:
+	driverLdconfigPath := "@" + filepath.Join(nvidiaDriverDir, strings.TrimPrefix(ldconfigPath, "@/"))
+
+	config.SetPath(nvidiaContainerCliKey("root"), nvidiaDriverDir)
+	config.SetPath(nvidiaContainerCliKey("path"), nvidiaContainerCliExecutablePath)
+	config.SetPath(nvidiaContainerCliKey("ldconfig"), driverLdconfigPath)
+
+	// Set the debug options if selected
+	debugOptions := map[string]string{
+		"nvidia-container-runtime.debug":     nvidiaContainerRuntimeDebugFlag,
+		"nvidia-container-runtime.log-level": nvidiaContainerRuntimeLogLevelFlag,
+		"nvidia-container-cli.debug":         nvidiaContainerCLIDebugFlag,
+	}
+	for key, value := range debugOptions {
+		if value == "" {
+			continue
+		}
+		if config.Get(key) != nil {
+			continue
+		}
+		config.Set(key, value)
+	}
+
+	_, err = config.WriteTo(targetConfig)
+	if err != nil {
+		return fmt.Errorf("error writing config: %v", err)
+	}
+	return nil
+}
+
+// installContainerCLI sets up the NVIDIA container CLI executable, copying the executable
+// and implementing the required wrapper
+func installContainerCLI(toolkitDir string) (string, error) {
+	log.Infof("Installing NVIDIA container CLI from '%v'", nvidiaContainerCliSource)
+
+	env := map[string]string{
+		"LD_LIBRARY_PATH": toolkitDir,
+	}
+
+	e := executable{
+		source: nvidiaContainerCliSource,
+		target: executableTarget{
+			dotfileName: "nvidia-container-cli.real",
+			wrapperName: "nvidia-container-cli",
+		},
+		env: env,
+	}
+
+	installedPath, err := e.install(toolkitDir)
+	if err != nil {
+		return "", fmt.Errorf("error installing NVIDIA container CLI: %v", err)
+	}
+	return installedPath, nil
+}
+
+// installRuntimeHook sets up the NVIDIA runtime hook, copying the executable
+// and implementing the required wrapper
+func installRuntimeHook(toolkitDir string, configFilePath string) (string, error) {
+	log.Infof("Installing NVIDIA container runtime hook from '%v'", nvidiaContainerRuntimeHookSource)
+
+	argLines := []string{
+		fmt.Sprintf("-config \"%s\"", configFilePath),
+	}
+
+	e := executable{
+		source: nvidiaContainerRuntimeHookSource,
+		target: executableTarget{
+			dotfileName: "nvidia-container-toolkit.real",
+			wrapperName: "nvidia-container-toolkit",
+		},
+		argLines: argLines,
+	}
+
+	installedPath, err := e.install(toolkitDir)
+	if err != nil {
+		return "", fmt.Errorf("error installing NVIDIA container runtime hook: %v", err)
+	}
+
+	err = installSymlink(toolkitDir, "nvidia-container-runtime-hook", installedPath)
+	if err != nil {
+		return "", fmt.Errorf("error installing symlink to NVIDIA container runtime hook: %v", err)
+	}
+
+	return installedPath, nil
+}
+
+// installSymlink creates a symlink in the toolkitDirectory that points to the specified target.
+// Note: The target is assumed to be local to the toolkit directory
+func installSymlink(toolkitDir string, link string, target string) error {
+	symlinkPath := filepath.Join(toolkitDir, link)
+	targetPath := filepath.Base(target)
+	log.Infof("Creating symlink '%v' -> '%v'", symlinkPath, targetPath)
+
+	err := os.Symlink(targetPath, symlinkPath)
+	if err != nil {
+		return fmt.Errorf("error creating symlink '%v' => '%v': %v", symlinkPath, targetPath, err)
+	}
+	return nil
+}
+
+// installFileToFolder copies a source file to a destination folder.
+// The path of the input file is ignored.
+// e.g. installFileToFolder("/some/path/file.txt", "/output/path")
+// will result in a file "/output/path/file.txt" being generated
+func installFileToFolder(destFolder string, src string) (string, error) {
+	name := filepath.Base(src)
+	return installFileToFolderWithName(destFolder, name, src)
+}
+
+// cp src destFolder/name
+func installFileToFolderWithName(destFolder string, name, src string) (string, error) {
+	dest := filepath.Join(destFolder, name)
+	err := installFile(dest, src)
+	if err != nil {
+		return "", fmt.Errorf("error copying '%v' to '%v': %v", src, dest, err)
+	}
+	return dest, nil
+}
+
+// installFile copies a file from src to dest and maintains
+// file modes
+func installFile(dest string, src string) error {
+	log.Infof("Installing '%v' to '%v'", src, dest)
+
+	source, err := os.Open(src)
+	if err != nil {
+		return fmt.Errorf("error opening source: %v", err)
+	}
+	defer source.Close()
+
+	destination, err := os.Create(dest)
+	if err != nil {
+		return fmt.Errorf("error creating destination: %v", err)
+	}
+	defer destination.Close()
+
+	_, err = io.Copy(destination, source)
+	if err != nil {
+		return fmt.Errorf("error copying file: %v", err)
+	}
+
+	err = applyModeFromSource(dest, src)
+	if err != nil {
+		return fmt.Errorf("error setting destination file mode: %v", err)
+	}
+	return nil
+}
+
+// applyModeFromSource sets the file mode for a destination file
+// to match that of a specified source file
+func applyModeFromSource(dest string, src string) error {
+	sourceInfo, err := os.Stat(src)
+	if err != nil {
+		return fmt.Errorf("error getting file info for '%v': %v", src, err)
+	}
+	err = os.Chmod(dest, sourceInfo.Mode())
+	if err != nil {
+		return fmt.Errorf("error setting mode for '%v': %v", dest, err)
+	}
+	return nil
+}
+
+// findLibrary searches a set of candidate libraries in the specified root for
+// a given library name
+func findLibrary(root string, libName string) (string, error) {
+	log.Infof("Finding library %v (root=%v)", libName, root)
+
+	candidateDirs := []string{
+		"/usr/lib64",
+		"/usr/lib/x86_64-linux-gnu",
+	}
+
+	for _, d := range candidateDirs {
+		l := filepath.Join(root, d, libName)
+		log.Infof("Checking library candidate '%v'", l)
+
+		libraryCandidate, err := resolveLink(l)
+		if err != nil {
+			log.Infof("Skipping library candidate '%v': %v", l, err)
+			continue
+		}
+
+		return libraryCandidate, nil
+	}
+
+	return "", fmt.Errorf("error locating library '%v'", libName)
+}
+
+// resolveLink finds the target of a symlink or the file itself in the
+// case of a regular file.
+// This is equivalent to running `readlink -f ${l}`
+func resolveLink(l string) (string, error) {
+	resolved, err := filepath.EvalSymlinks(l)
+	if err != nil {
+		return "", fmt.Errorf("error resolving link '%v': %v", l, err)
+	}
+	if l != resolved {
+		log.Infof("Resolved link: '%v' => '%v'", l, resolved)
+	}
+	return resolved, nil
+}
+
+func createDirectories(dir ...string) error {
+	for _, d := range dir {
+		log.Infof("Creating directory '%v'", d)
+		err := os.MkdirAll(d, 0755)
+		if err != nil {
+			return fmt.Errorf("error creating directory: %v", err)
+		}
+	}
+	return nil
+}