Skip to content

Commit

Permalink
use the new Go wrapper program
Browse files Browse the repository at this point in the history 
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 <[email protected]>
  • Loading branch information
@jfroy
jfroy committed Nov 6, 2024
1 parent dc79fa8 commit 8d8dbd3
Show file tree
Hide file tree
Showing 5 changed files with 171 additions and 180 deletions.
140 changes: 75 additions & 65 deletions tools/container/toolkit/executable.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,123 +18,133 @@ 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)
}
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 <PATH, which instructs wrapper to insert the value at the head of a
// colon-separated environment variable list.
delete(envm, "PATH")
envm["<PATH"] = destFolder
}
// 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))
var envv []string
for k, v := range envm {
envv = append(envv, k+"="+r.apply(v))
}
sort.Strings(envv)
for _, e := range envv {
fmt.Fprintf(f, "%s\n", e)
}
// Add the script arguments "$@"
fmt.Fprintln(wrapper, "\t\"$@\"")

return nil
}

Expand Down
139 changes: 83 additions & 56 deletions tools/container/toolkit/executable_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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{
"PATH": "some-path",
source: "source",
envm: map[string]string{
"FOO": "BAR",
},
},
expectedLines: []string{
shebang,
"PATH=/dest/folder:some-path \\",
"source.real \\",
"\t\"$@\"",
"",
expectedEnvv: []string{
fmt.Sprintf("<PATH=%s", destDirPattern),
"FOO=BAR",
},
},
{
e: executable{
preLines: []string{
"preline1",
"preline2",
source: "source",
envm: map[string]string{
"PATH": "some-path",
"FOO": "BAR",
},
},
expectedLines: []string{
shebang,
"preline1",
"preline2",
"PATH=/dest/folder:$PATH \\",
"source.real \\",
"\t\"$@\"",
"",
expectedEnvv: []string{
"FOO=BAR",
fmt.Sprintf("PATH=%s:some-path", destDirPattern),
},
},
{
e: executable{
argLines: []string{
"argline1",
"argline2",
source: "source",
argv: []string{
"argb",
"arga",
"argc",
},
},
expectedLines: []string{
shebang,
"PATH=/dest/folder:$PATH \\",
"source.real \\",
"\targline1 \\",
"\targline2 \\",
"\t\"$@\"",
"",
expectedArgv: []string{
"argb",
"arga",
"argc",
},
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)

exepectedContents := strings.Join(tc.expectedLines, "\n")
require.Equal(t, exepectedContents, buf.String(), "%v: %v", i, tc)
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)

}
}
}

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)
Expand All @@ -123,7 +123,6 @@ func TestInstallExecutable(t *testing.T) {
e := executable{
source: source,
target: executableTarget{
dotfileName: "input.real",
wrapperName: "input",
},
}
Expand All @@ -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
}
Loading

0 comments on commit 8d8dbd3

Please sign in to comment.