Merge d9c52ecd4e into 178369eb8e

tools/container/toolkit/executable.go

@@ -18,71 +18,78 @@ package toolkit
 
 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
+	argv   []string
-	preLines []string
+	envm   map[string]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) {
+	if destFolder == "" {
+		return "", fmt.Errorf("destination folder must be specified")
+	}
+	if e.source == "" {
+		return "", fmt.Errorf("source executable must be specified")
+	}
 	log.Infof("Installing executable '%v' to %v", e.source, destFolder)
-
+	dotRealFilename := e.dotRealFilename()
-	dotfileName := e.dotfileName()
+	dotRealPath, err := installFileToFolderWithName(destFolder, dotRealFilename, e.source)
-
-	installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source)
 	if err != nil {
-		return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotfileName, err)
+		return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotRealFilename, err)
 	}
-	log.Infof("Installed '%v'", installedDotfileName)
+	log.Infof("Installed '%v'", dotRealPath)
 
-	wrapperFilename, err := e.installWrapper(destFolder, installedDotfileName)
+	wrapperPath, err := e.installWrapper(destFolder)
 	if err != nil {
-		return "", fmt.Errorf("error wrapping '%v': %v", installedDotfileName, err)
+		return "", fmt.Errorf("error installing wrapper: %v", err)
 	}
-	log.Infof("Installed wrapper '%v'", wrapperFilename)
+	log.Infof("Installed wrapper '%v'", wrapperPath)
-
+	return wrapperPath, nil
-	return wrapperFilename, nil
 }
 
-func (e executable) dotfileName() string {
+func (e executable) dotRealFilename() string {
-	return e.target.dotfileName
+	return e.wrapperName() + ".real"
 }
 
 func (e executable) wrapperName() string {
+	if e.target.wrapperName == "" {
+		return filepath.Base(e.source)
+	}
 	return e.target.wrapperName
 }
 
-func (e executable) installWrapper(destFolder string, dotfileName string) (string, error) {
+func (e executable) installWrapper(destFolder string) (string, error) {
-	wrapperPath := filepath.Join(destFolder, e.wrapperName())
+	currentExe, err := os.Executable()
-	wrapper, err := os.Create(wrapperPath)
 	if err != nil {
-		return "", fmt.Errorf("error creating executable wrapper: %v", err)
+		return "", fmt.Errorf("error getting current executable: %v", err)
 	}
-	defer wrapper.Close()
+	src := filepath.Join(filepath.Dir(currentExe), "wrapper")
-
+	wrapperPath, err := installFileToFolderWithName(destFolder, e.wrapperName(), src)
-	err = e.writeWrapperTo(wrapper, destFolder, dotfileName)
 	if err != nil {
-		return "", fmt.Errorf("error writing wrapper contents: %v", err)
+		return "", fmt.Errorf("error installing wrapper program: %v", err)
+	}
+	err = e.writeWrapperArgv(wrapperPath, destFolder)
+	if err != nil {
+		return "", fmt.Errorf("error writing wrapper argv: %v", err)
+	}
+	err = e.writeWrapperEnvv(wrapperPath, destFolder)
+	if err != nil {
+		return "", fmt.Errorf("error writing wrapper envv: %v", err)
 	}
-
 	err = ensureExecutable(wrapperPath)
 	if err != nil {
 		return "", fmt.Errorf("error making wrapper executable: %v", err)
@@ -90,51 +97,54 @@ func (e executable) installWrapper(destFolder string, dotfileName string) (strin
 	return wrapperPath, nil
 }
 
-func (e executable) writeWrapperTo(wrapper io.Writer, destFolder string, dotfileName string) error {
+func (e executable) writeWrapperArgv(wrapperPath string, destFolder string) error {
+	if e.argv == nil {
+		return nil
+	}
 	r := newReplacements(destDirPattern, destFolder)
-
+	f, err := os.OpenFile(wrapperPath+".argv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
-	// Add the shebang
+	if err != nil {
-	fmt.Fprintln(wrapper, "#! /bin/sh")
+		return err
-
+	}
-	// Add the preceding lines if any
+	defer f.Close()
-	for _, line := range e.preLines {
+	for _, arg := range e.argv {
-		fmt.Fprintf(wrapper, "%s\n", r.apply(line))
+		fmt.Fprintf(f, "%s\n", r.apply(arg))
+	}
+	return nil
 }
 
-	// Update the path to include the destination folder
+func (e executable) writeWrapperEnvv(wrapperPath string, destFolder string) error {
-	var env map[string]string
+	r := newReplacements(destDirPattern, destFolder)
-	if e.env == nil {
+	f, err := os.OpenFile(wrapperPath+".envv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
-		env = make(map[string]string)
+	if err != nil {
+		return err
+	}
+	defer f.Close()
+
+	// Update PATH to insert the destination folder at the head.
+	var envm map[string]string
+	if e.envm == nil {
+		envm = make(map[string]string)
 	} else {
-		env = e.env
+		envm = e.envm
+	}
+	if path, ok := envm["PATH"]; ok {
+		envm["PATH"] = destFolder + ":" + path
+	} else {
+		// Replace PATH with <PATH, which instructs wrapper to insert the value at the head of a
+		// colon-separated environment variable list.
+		delete(envm, "PATH")
+		envm["<PATH"] = destFolder
 	}
 
-	path, specified := env["PATH"]
+	var envv []string
-	if !specified {
+	for k, v := range envm {
-		path = "$PATH"
+		envv = append(envv, k+"="+r.apply(v))
 	}
-	env["PATH"] = strings.Join([]string{destFolder, path}, ":")
+	sort.Strings(envv)
-
+	for _, e := range envv {
-	var sortedEnvvars []string
+		fmt.Fprintf(f, "%s\n", e)
-	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
 }
 

tools/container/toolkit/executable_test.go

@@ -17,102 +17,102 @@
 package toolkit
 
 import (
-	"bytes"
+	"bufio"
+	"fmt"
+	"io/fs"
 	"os"
 	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
 )
 
 func TestWrapper(t *testing.T) {
-	const shebang = "#! /bin/sh"
+	createTestWrapperProgram(t)
-	const destFolder = "/dest/folder"
-	const dotfileName = "source.real"
 
 	testCases := []struct {
 		e            executable
-		expectedLines []string
+		expectedArgv []string
+		expectedEnvv []string
 	}{
 		{
-			e: executable{},
+			e: executable{source: "source"},
-			expectedLines: []string{
+			expectedEnvv: []string{
-				shebang,
+				fmt.Sprintf("<PATH=%s", destDirPattern),
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
 			},
 		},
 		{
 			e: executable{
-				env: map[string]string{
+				source: "source",
+				envm: map[string]string{
+					"FOO": "BAR",
+				},
+			},
+			expectedEnvv: []string{
+				fmt.Sprintf("<PATH=%s", destDirPattern),
+				"FOO=BAR",
+			},
+		},
+		{
+			e: executable{
+				source: "source",
+				envm: map[string]string{
 					"PATH": "some-path",
+					"FOO":  "BAR",
 				},
 			},
-			expectedLines: []string{
+			expectedEnvv: []string{
-				shebang,
+				"FOO=BAR",
-				"PATH=/dest/folder:some-path \\",
+				fmt.Sprintf("PATH=%s:some-path", destDirPattern),
-				"source.real \\",
-				"\t\"$@\"",
-				"",
 			},
 		},
 		{
 			e: executable{
-				preLines: []string{
+				source: "source",
-					"preline1",
+				argv: []string{
-					"preline2",
+					"argb",
+					"arga",
+					"argc",
 				},
 			},
-			expectedLines: []string{
+			expectedArgv: []string{
-				shebang,
+				"argb",
-				"preline1",
+				"arga",
-				"preline2",
+				"argc",
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
 			},
-		},
+			expectedEnvv: []string{
-		{
+				fmt.Sprintf("<PATH=%s", destDirPattern),
-			e: executable{
-				argLines: []string{
-					"argline1",
-					"argline2",
-				},
-			},
-			expectedLines: []string{
-				shebang,
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\targline1 \\",
-				"\targline2 \\",
-				"\t\"$@\"",
-				"",
 			},
 		},
 	}
 
-	for i, tc := range testCases {
+	for _, tc := range testCases {
-		buf := &bytes.Buffer{}
+		destFolder := t.TempDir()
-
+		r := newReplacements(destDirPattern, destFolder)
-		err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
+		for k, v := range tc.expectedEnvv {
+			tc.expectedEnvv[k] = r.apply(v)
+		}
+		path, err := tc.e.installWrapper(destFolder)
 		require.NoError(t, err)
+		require.FileExists(t, path)
+		envv, err := readAllLines(path + ".envv")
+		require.NoError(t, err)
+		require.Equal(t, tc.expectedEnvv, envv)
+		argv, err := readAllLines(path + ".argv")
+		if tc.expectedArgv == nil {
+			require.ErrorAs(t, err, &fs.ErrNotExist)
+		} else {
+			require.Equal(t, tc.expectedArgv, argv)
 
-		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("", "")
+	createTestWrapperProgram(t)
-	require.NoError(t, err)
-	defer os.RemoveAll(inputFolder)
 
 	// Create the source file
-	source := filepath.Join(inputFolder, "input")
+	source := filepath.Join(t.TempDir(), "input")
 	sourceFile, err := os.Create(source)
 
 	base := filepath.Base(source)
@@ -123,7 +123,6 @@ func TestInstallExecutable(t *testing.T) {
 	e := executable{
 		source: source,
 		target: executableTarget{
-			dotfileName: "input.real",
 			wrapperName: "input",
 		},
 	}
@@ -150,3 +149,31 @@ func TestInstallExecutable(t *testing.T) {
 	require.NoError(t, err)
 	require.NotEqual(t, 0, wrapperInfo.Mode()&0111)
 }
+
+func createTestWrapperProgram(t *testing.T) {
+	t.Helper()
+	currentExe, err := os.Executable()
+	if err != nil {
+		t.Fatalf("error getting current executable: %v", err)
+	}
+	wrapperPath := filepath.Join(filepath.Dir(currentExe), "wrapper")
+	f, err := os.Create(wrapperPath)
+	if err != nil {
+		t.Fatalf("error creating test wrapper: %v", err)
+	}
+	f.Close()
+}
+
+func readAllLines(path string) (s []string, err error) {
+	f, err := os.Open(path)
+	if err != nil {
+		return
+	}
+	defer f.Close()
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		s = append(s, scanner.Text())
+	}
+	err = scanner.Err()
+	return
+}

tools/container/toolkit/runtime.go

@@ -48,37 +48,22 @@ func installContainerRuntimes(sourceRoot string, toolkitDir string) error {
 // created to allow for the configuration of the runtime environment.
 func newNvidiaContainerRuntimeInstaller(source string) *executable {
 	wrapperName := filepath.Base(source)
-	dotfileName := wrapperName + ".real"
 	target := executableTarget{
-		dotfileName: dotfileName,
 		wrapperName: wrapperName,
 	}
 	return newRuntimeInstaller(source, target, nil)
 }
 
 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,
+		envm:   runtimeEnv,
-		preLines: preLines,
 	}
-
 	return &r
 }

tools/container/toolkit/runtime_test.go

@@ -17,8 +17,7 @@
 package toolkit
 
 import (
-	"bytes"
+	"path/filepath"
-	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -26,32 +25,10 @@ import (
 
 func TestNvidiaContainerRuntimeInstallerWrapper(t *testing.T) {
 	r := newNvidiaContainerRuntimeInstaller(nvidiaContainerRuntimeSource)
-
+	require.Equal(t, nvidiaContainerRuntimeSource, r.source)
-	const shebang = "#! /bin/sh"
+	require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource), r.target.wrapperName)
-	const destFolder = "/dest/folder"
+	require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource), r.wrapperName())
-	const dotfileName = "source.real"
+	require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource)+".real", r.dotRealFilename())
-
+	require.Nil(t, r.argv)
-	buf := &bytes.Buffer{}
+	require.Equal(t, map[string]string{"XDG_CONFIG_HOME": filepath.Join(destDirPattern, ".config")}, r.envm)
-
-	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())
 }

tools/container/toolkit/toolkit.go

@@ -529,7 +529,6 @@ func installContainerToolkitCLI(sourceRoot string, toolkitDir string) (string, e
 	e := executable{
 		source: filepath.Join(sourceRoot, "/usr/bin/nvidia-ctk"),
 		target: executableTarget{
-			dotfileName: "nvidia-ctk.real",
 			wrapperName: "nvidia-ctk",
 		},
 	}
@@ -542,7 +541,6 @@ func installContainerCDIHookCLI(sourceRoot string, toolkitDir string) (string, e
 	e := executable{
 		source: filepath.Join(sourceRoot, "/usr/bin/nvidia-cdi-hook"),
 		target: executableTarget{
-			dotfileName: "nvidia-cdi-hook.real",
 			wrapperName: "nvidia-cdi-hook",
 		},
 	}
@@ -555,17 +553,16 @@ func installContainerCDIHookCLI(sourceRoot string, toolkitDir string) (string, e
 func installContainerCLI(sourceRoot string, toolkitRoot string) (string, error) {
 	log.Infof("Installing NVIDIA container CLI from '%v'", nvidiaContainerCliSource)
 
-	env := map[string]string{
+	envm := map[string]string{
 		"LD_LIBRARY_PATH": toolkitRoot,
 	}
 
 	e := executable{
 		source: filepath.Join(sourceRoot, nvidiaContainerCliSource),
 		target: executableTarget{
-			dotfileName: "nvidia-container-cli.real",
 			wrapperName: "nvidia-container-cli",
 		},
-		env: env,
+		envm: envm,
 	}
 
 	installedPath, err := e.install(toolkitRoot)
@@ -580,17 +577,12 @@ func installContainerCLI(sourceRoot string, toolkitRoot string) (string, error)
 func installRuntimeHook(sourceRoot string, toolkitRoot 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: filepath.Join(sourceRoot, nvidiaContainerRuntimeHookSource),
 		target: executableTarget{
-			dotfileName: "nvidia-container-runtime-hook.real",
 			wrapperName: "nvidia-container-runtime-hook",
 		},
-		argLines: argLines,
+		argv: []string{"-config", configFilePath},
 	}
 
 	installedPath, err := e.install(toolkitRoot)

tools/container/wrapper/wrapper.go

@@ -0,0 +1,89 @@
+/**
+# Copyright (c) 2024, 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 (
+	"bufio"
+	"errors"
+	"io/fs"
+	"log"
+	"os"
+	"strings"
+
+	"golang.org/x/sys/unix"
+)
+
+func main() {
+	program, err := os.Executable()
+	if err != nil {
+		log.Fatalf("failed to get executable: %v", err)
+	}
+	argv := makeArgv(program)
+	envv := makeEnvv(program)
+	if err := unix.Exec(program+".real", argv, envv); err != nil {
+		log.Fatalf("failed to exec %s: %v", program+".real", err)
+	}
+
+}
+
+func makeArgv(program string) []string {
+	argv := []string{os.Args[0] + ".real"}
+	f, err := os.Open(program + ".argv")
+	if err != nil {
+		if !errors.Is(err, fs.ErrNotExist) {
+			log.Printf("failed to open argv file: %v", err)
+		}
+		return append(argv, os.Args[1:]...)
+	}
+	defer f.Close()
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		argv = append(argv, scanner.Text())
+	}
+	if err := scanner.Err(); err != nil {
+		log.Fatalf("failed to read argv file: %v", err)
+	}
+	return append(argv, os.Args[1:]...)
+}
+
+func makeEnvv(program string) []string {
+	f, err := os.Open(program + ".envv")
+	if err != nil {
+		if !errors.Is(err, fs.ErrNotExist) {
+			log.Printf("failed to open env file: %v", err)
+		}
+		return os.Environ()
+	}
+	defer f.Close()
+	var env []string
+	scanner := bufio.NewScanner(f)
+	for scanner.Scan() {
+		kv := strings.SplitN(scanner.Text(), "=", 2)
+		if strings.HasPrefix(kv[0], "<") {
+			kv[0] = kv[0][1:]
+			kv[1] = kv[1] + ":" + os.Getenv(kv[0])
+		} else if strings.HasPrefix(kv[0], ">") {
+			kv[0] = kv[0][1:]
+			kv[1] = os.Getenv(kv[0]) + ":" + kv[1]
+		}
+		env = append(env, kv[0]+"="+kv[1])
+	}
+	if err := scanner.Err(); err != nil {
+		log.Fatalf("failed to read argv file: %v", err)
+	}
+	return append(env, os.Environ()...)
+}