From dc79fa851388737a3ff075117caa3256b3ba0eec Mon Sep 17 00:00:00 2001 From: Jean-Francois Roy Date: Wed, 18 Sep 2024 08:47:21 -0700 Subject: [PATCH 1/4] implement golang wrapper to replace shell scripts Some platforms and Kubernetes distributions do not include a shell. This patch replaces the shell wrapper scripts with a small Go program. Signed-off-by: Jean-Francois Roy --- tools/container/wrapper/wrapper.go | 115 +++++++++++++++++++++++++++++ 1 file changed, 115 insertions(+) create mode 100644 tools/container/wrapper/wrapper.go diff --git a/tools/container/wrapper/wrapper.go b/tools/container/wrapper/wrapper.go new file mode 100644 index 000000000..9c9761aa0 --- /dev/null +++ b/tools/container/wrapper/wrapper.go @@ -0,0 +1,115 @@ +/** +# 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" + "os/exec" + "path/filepath" + "strings" + + "golang.org/x/sys/unix" +) + +func main() { + program, err := os.Executable() + if err != nil { + log.Fatalf("failed to get executable: %v", err) + } + if isRuntimeWrapper(program) && !isNvidiaModuleLoaded() { + log.Println("nvidia driver modules are not yet loaded, invoking runc directly") + program, err := exec.LookPath("runc") + if err != nil { + log.Fatalf("failed to find runc: %v", err) + } + argv := []string{"runc"} + argv = append(argv, os.Args[1:]...) + execve(program, argv, os.Environ()) + } + argv := makeArgv(program) + envv := makeEnvv(program) + execve(program+".real", argv, envv) +} + +func isRuntimeWrapper(program string) bool { + return filepath.Base(program) == "nvidia-container-runtime" || + filepath.Base(program) == "nvidia-container-runtime.cdi" || + filepath.Base(program) == "nvidia-container-runtime.legacy" +} + +func isNvidiaModuleLoaded() bool { + _, err := os.Stat("/proc/driver/nvidia/version") + return err == nil +} + +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()...) +} + +func execve(program string, argv []string, envv []string) { + if err := unix.Exec(program, argv, envv); err != nil { + log.Fatalf("failed to exec %s: %v", program, err) + } +} From 8d8dbd38c3353addf1599e268e0a36852e84eb83 Mon Sep 17 00:00:00 2001 From: Jean-Francois Roy Date: Wed, 18 Sep 2024 08:49:38 -0700 Subject: [PATCH 2/4] use the new Go wrapper program This patch modifies the the container toolkit installer, used by the GPU operator, to use the new Go wrapper program. Signed-off-by: Jean-Francois Roy --- tools/container/toolkit/executable.go | 140 +++++++++++---------- tools/container/toolkit/executable_test.go | 139 +++++++++++--------- tools/container/toolkit/runtime.go | 21 +--- tools/container/toolkit/runtime_test.go | 37 ++---- tools/container/toolkit/toolkit.go | 14 +-- 5 files changed, 171 insertions(+), 180 deletions(-) diff --git a/tools/container/toolkit/executable.go b/tools/container/toolkit/executable.go index 394ca0076..1b8c4f1c5 100644 --- a/tools/container/toolkit/executable.go +++ b/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 - preLines []string - argLines []string + source string + target executableTarget + argv []string + envm map[string]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) - - dotfileName := e.dotfileName() - - installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source) + dotRealFilename := e.dotRealFilename() + dotRealPath, err := installFileToFolderWithName(destFolder, dotRealFilename, 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) - - return wrapperFilename, nil + log.Infof("Installed wrapper '%v'", wrapperPath) + return wrapperPath, nil } -func (e executable) dotfileName() string { - return e.target.dotfileName +func (e executable) dotRealFilename() string { + 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) { - wrapperPath := filepath.Join(destFolder, e.wrapperName()) - wrapper, err := os.Create(wrapperPath) +func (e executable) installWrapper(destFolder string) (string, error) { + currentExe, err := os.Executable() if err != nil { - return "", fmt.Errorf("error creating executable wrapper: %v", err) + return "", fmt.Errorf("error getting current executable: %v", err) } - defer wrapper.Close() - - err = e.writeWrapperTo(wrapper, destFolder, dotfileName) + src := filepath.Join(filepath.Dir(currentExe), "wrapper") + wrapperPath, err := installFileToFolderWithName(destFolder, e.wrapperName(), src) 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) - - // 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)) + f, err := os.OpenFile(wrapperPath+".argv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440) + if err != nil { + return err } - - // 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 + defer f.Close() + for _, arg := range e.argv { + fmt.Fprintf(f, "%s\n", r.apply(arg)) } + return nil +} - path, specified := env["PATH"] - if !specified { - path = "$PATH" +func (e executable) writeWrapperEnvv(wrapperPath string, destFolder string) error { + r := newReplacements(destDirPattern, destFolder) + f, err := os.OpenFile(wrapperPath+".envv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440) + if err != nil { + return err } - env["PATH"] = strings.Join([]string{destFolder, path}, ":") + defer f.Close() - var sortedEnvvars []string - for e := range env { - sortedEnvvars = append(sortedEnvvars, e) + // 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 { + envm = e.envm } - sort.Strings(sortedEnvvars) - - for _, e := range sortedEnvvars { - v := env[e] - fmt.Fprintf(wrapper, "%s=%s \\\n", e, r.apply(v)) + if path, ok := envm["PATH"]; ok { + envm["PATH"] = destFolder + ":" + path + } else { + // Replace PATH with /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, + source: source, + target: target, + envm: runtimeEnv, } - return &r } diff --git a/tools/container/toolkit/runtime_test.go b/tools/container/toolkit/runtime_test.go index d2841506d..f8b411781 100644 --- a/tools/container/toolkit/runtime_test.go +++ b/tools/container/toolkit/runtime_test.go @@ -17,8 +17,7 @@ package toolkit import ( - "bytes" - "strings" + "path/filepath" "testing" "github.com/stretchr/testify/require" @@ -26,32 +25,10 @@ import ( func TestNvidiaContainerRuntimeInstallerWrapper(t *testing.T) { r := newNvidiaContainerRuntimeInstaller(nvidiaContainerRuntimeSource) - - 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()) + require.Equal(t, nvidiaContainerRuntimeSource, r.source) + require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource), r.target.wrapperName) + require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource), r.wrapperName()) + require.Equal(t, filepath.Base(nvidiaContainerRuntimeSource)+".real", r.dotRealFilename()) + require.Nil(t, r.argv) + require.Equal(t, map[string]string{"XDG_CONFIG_HOME": filepath.Join(destDirPattern, ".config")}, r.envm) } diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go index 9b97b4197..5e6a3b8b2 100644 --- a/tools/container/toolkit/toolkit.go +++ b/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) From 0c1e76a22103fed5fcde449a1aa2bde9bdda7d32 Mon Sep 17 00:00:00 2001 From: Jean-Francois Roy Date: Wed, 30 Oct 2024 08:57:21 -0700 Subject: [PATCH 3/4] inline execve Signed-off-by: Jean-Francois Roy --- tools/container/wrapper/wrapper.go | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/tools/container/wrapper/wrapper.go b/tools/container/wrapper/wrapper.go index 9c9761aa0..ac6ef90f5 100644 --- a/tools/container/wrapper/wrapper.go +++ b/tools/container/wrapper/wrapper.go @@ -42,11 +42,16 @@ func main() { } argv := []string{"runc"} argv = append(argv, os.Args[1:]...) - execve(program, argv, os.Environ()) + if err := unix.Exec(program, argv, os.Environ()); err != nil { + log.Fatalf("failed to exec %s: %v", program, err) + } } argv := makeArgv(program) envv := makeEnvv(program) - execve(program+".real", argv, envv) + if err := unix.Exec(program+".real", argv, envv); err != nil { + log.Fatalf("failed to exec %s: %v", program+".real", err) + } + } func isRuntimeWrapper(program string) bool { @@ -107,9 +112,3 @@ func makeEnvv(program string) []string { } return append(env, os.Environ()...) } - -func execve(program string, argv []string, envv []string) { - if err := unix.Exec(program, argv, envv); err != nil { - log.Fatalf("failed to exec %s: %v", program, err) - } -} From d9c52ecd4e922eeaebd9e20fcbb0a11a0330157d Mon Sep 17 00:00:00 2001 From: Jean-Francois Roy Date: Thu, 31 Oct 2024 09:18:39 -0700 Subject: [PATCH 4/4] remove runtime wrapper driver detection This will be handled by the runtime itself. Signed-off-by: Jean-Francois Roy --- tools/container/wrapper/wrapper.go | 25 ------------------------- 1 file changed, 25 deletions(-) diff --git a/tools/container/wrapper/wrapper.go b/tools/container/wrapper/wrapper.go index ac6ef90f5..93763678c 100644 --- a/tools/container/wrapper/wrapper.go +++ b/tools/container/wrapper/wrapper.go @@ -22,8 +22,6 @@ import ( "io/fs" "log" "os" - "os/exec" - "path/filepath" "strings" "golang.org/x/sys/unix" @@ -34,18 +32,6 @@ func main() { if err != nil { log.Fatalf("failed to get executable: %v", err) } - if isRuntimeWrapper(program) && !isNvidiaModuleLoaded() { - log.Println("nvidia driver modules are not yet loaded, invoking runc directly") - program, err := exec.LookPath("runc") - if err != nil { - log.Fatalf("failed to find runc: %v", err) - } - argv := []string{"runc"} - argv = append(argv, os.Args[1:]...) - if err := unix.Exec(program, argv, os.Environ()); err != nil { - log.Fatalf("failed to exec %s: %v", program, err) - } - } argv := makeArgv(program) envv := makeEnvv(program) if err := unix.Exec(program+".real", argv, envv); err != nil { @@ -54,17 +40,6 @@ func main() { } -func isRuntimeWrapper(program string) bool { - return filepath.Base(program) == "nvidia-container-runtime" || - filepath.Base(program) == "nvidia-container-runtime.cdi" || - filepath.Base(program) == "nvidia-container-runtime.legacy" -} - -func isNvidiaModuleLoaded() bool { - _, err := os.Stat("/proc/driver/nvidia/version") - return err == nil -} - func makeArgv(program string) []string { argv := []string{os.Args[0] + ".real"} f, err := os.Open(program + ".argv")