UniqAI/nvidia-container-toolkit
1
0
Fork 0
mirror of https://github.com/NVIDIA/nvidia-container-toolkit synced 2025-04-03 20:30:48 +00:00
Code Issues Actions 4 Packages Projects Releases Wiki Activity

Merge d9c52ecd4e into 178369eb8e

Browse Source
This commit is contained in:
Jean-François Roy 2025-01-08 09:35:02 -08:00 committed by GitHub
commit 3cc5eaa53f
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 265 additions and 185 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

144
tools/container/toolkit/executable.go
View File

@ -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
}
defer f.Close()
for _, arg := range e.argv {
fmt.Fprintf(f, "%s\n", r.apply(arg))
}
return nil
}
// Update the path to include the destination folder
var env map[string]string
if e.env == nil {
env = make(map[string]string)
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
}
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"]
if !specified {
path = "$PATH"
var envv []string
for k, v := range envm {
envv = append(envv, k+"="+r.apply(v))
}
env["PATH"] = strings.Join([]string{destFolder, path}, ":")
var sortedEnvvars []string
for e := range env {
sortedEnvvars = append(sortedEnvvars, e)
sort.Strings(envv)
for _, e := range envv {
fmt.Fprintf(f, "%s\n", 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
}

145
tools/container/toolkit/executable_test.go
View File

@ -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"
const destFolder = "/dest/folder"
const dotfileName = "source.real"
createTestWrapperProgram(t)
testCases := []struct {
e executable
expectedLines []string
e executable
expectedArgv []string
expectedEnvv []string
}{
{
e: executable{},
expectedLines: []string{
shebang,
"PATH=/dest/folder:$PATH \\",
"source.real \\",
"\t\"$@\"",
"",
e: executable{source: "source"},
expectedEnvv: []string{
fmt.Sprintf("<PATH=%s", destDirPattern),
},
},
{
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{
shebang,
"PATH=/dest/folder:some-path \\",
"source.real \\",
"\t\"$@\"",
"",
expectedEnvv: []string{
"FOO=BAR",
fmt.Sprintf("PATH=%s:some-path", destDirPattern),
},
},
{
e: executable{
preLines: []string{
"preline1",
"preline2",
source: "source",
argv: []string{
"argb",
"arga",
"argc",
},
},
expectedLines: []string{
shebang,
"preline1",
"preline2",
"PATH=/dest/folder:$PATH \\",
"source.real \\",
"\t\"$@\"",
"",
expectedArgv: []string{
"argb",
"arga",
"argc",
},
},
{
e: executable{
argLines: []string{
"argline1",
"argline2",
},
},
expectedLines: []string{
shebang,
"PATH=/dest/folder:$PATH \\",
"source.real \\",
"\targline1 \\",
"\targline2 \\",
"\t\"$@\"",
"",
expectedEnvv: []string{
fmt.Sprintf("<PATH=%s", destDirPattern),
},
},
}
for i, tc := range testCases {
buf := &bytes.Buffer{}
err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
for _, tc := range testCases {
destFolder := t.TempDir()
r := newReplacements(destDirPattern, destFolder)
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("", "")
require.NoError(t, err)
defer os.RemoveAll(inputFolder)
createTestWrapperProgram(t)
// 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
}

21
tools/container/toolkit/runtime.go
View File

@ -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,
preLines: preLines,
source: source,
target: target,
envm: runtimeEnv,
}
return &r
}

37
tools/container/toolkit/runtime_test.go
View File

@ -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)
}

14
tools/container/toolkit/toolkit.go
View File

@ -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)

89
tools/container/wrapper/wrapper.go Normal file
View File

@ -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()...)
}