Skip to content

Commit d291802

Browse files
jfroyjfroy
committed
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.
1 parent 20e4dca commit d291802

File tree

5 files changed

+171
-180
lines changed

5 files changed

+171
-180
lines changed

‎tools/container/toolkit/executable.go

+75-65
Original file line numberDiff line numberDiff line change
@@ -18,123 +18,133 @@ package main
1818

1919
import (
2020
"fmt"
21-
"io"
2221
"os"
2322
"path/filepath"
2423
"sort"
25-
"strings"
2624

2725
log "github.com/sirupsen/logrus"
2826
)
2927

3028
type executableTarget struct {
31-
dotfileName string
3229
wrapperName string
3330
}
3431

3532
type executable struct {
36-
source string
37-
target executableTarget
38-
env map[string]string
39-
preLines []string
40-
argLines []string
33+
source string
34+
target executableTarget
35+
argv []string
36+
envm map[string]string
4137
}
4238

4339
// install installs an executable component of the NVIDIA container toolkit. The source executable
4440
// is copied to a `.real` file and a wapper is created to set up the environment as required.
4541
func (e executable) install(destFolder string) (string, error) {
42+
if destFolder == "" {
43+
return "", fmt.Errorf("destination folder must be specified")
44+
}
45+
if e.source == "" {
46+
return "", fmt.Errorf("source executable must be specified")
47+
}
4648
log.Infof("Installing executable '%v' to %v", e.source, destFolder)
47-
48-
dotfileName := e.dotfileName()
49-
50-
installedDotfileName, err := installFileToFolderWithName(destFolder, dotfileName, e.source)
49+
dotRealFilename := e.dotRealFilename()
50+
dotRealPath, err := installFileToFolderWithName(destFolder, dotRealFilename, e.source)
5151
if err != nil {
52-
return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotfileName, err)
52+
return "", fmt.Errorf("error installing file '%v' as '%v': %v", e.source, dotRealFilename, err)
5353
}
54-
log.Infof("Installed '%v'", installedDotfileName)
54+
log.Infof("Installed '%v'", dotRealPath)
5555

56-
wrapperFilename, err := e.installWrapper(destFolder, installedDotfileName)
56+
wrapperPath, err := e.installWrapper(destFolder)
5757
if err != nil {
58-
return "", fmt.Errorf("error wrapping '%v': %v", installedDotfileName, err)
58+
return "", fmt.Errorf("error installing wrapper: %v", err)
5959
}
60-
log.Infof("Installed wrapper '%v'", wrapperFilename)
61-
62-
return wrapperFilename, nil
60+
log.Infof("Installed wrapper '%v'", wrapperPath)
61+
return wrapperPath, nil
6362
}
6463

65-
func (e executable) dotfileName() string {
66-
return e.target.dotfileName
64+
func (e executable) dotRealFilename() string {
65+
return e.wrapperName() + ".real"
6766
}
6867

6968
func (e executable) wrapperName() string {
69+
if e.target.wrapperName == "" {
70+
return filepath.Base(e.source)
71+
}
7072
return e.target.wrapperName
7173
}
7274

73-
func (e executable) installWrapper(destFolder string, dotfileName string) (string, error) {
74-
wrapperPath := filepath.Join(destFolder, e.wrapperName())
75-
wrapper, err := os.Create(wrapperPath)
75+
func (e executable) installWrapper(destFolder string) (string, error) {
76+
currentExe, err := os.Executable()
7677
if err != nil {
77-
return "", fmt.Errorf("error creating executable wrapper: %v", err)
78+
return "", fmt.Errorf("error getting current executable: %v", err)
7879
}
79-
defer wrapper.Close()
80-
81-
err = e.writeWrapperTo(wrapper, destFolder, dotfileName)
80+
src := filepath.Join(filepath.Dir(currentExe), "wrapper")
81+
wrapperPath, err := installFileToFolderWithName(destFolder, e.wrapperName(), src)
8282
if err != nil {
83-
return "", fmt.Errorf("error writing wrapper contents: %v", err)
83+
return "", fmt.Errorf("error installing wrapper program: %v", err)
84+
}
85+
err = e.writeWrapperArgv(wrapperPath, destFolder)
86+
if err != nil {
87+
return "", fmt.Errorf("error writing wrapper argv: %v", err)
88+
}
89+
err = e.writeWrapperEnvv(wrapperPath, destFolder)
90+
if err != nil {
91+
return "", fmt.Errorf("error writing wrapper envv: %v", err)
8492
}
85-
8693
err = ensureExecutable(wrapperPath)
8794
if err != nil {
8895
return "", fmt.Errorf("error making wrapper executable: %v", err)
8996
}
9097
return wrapperPath, nil
9198
}
9299

93-
func (e executable) writeWrapperTo(wrapper io.Writer, destFolder string, dotfileName string) error {
100+
func (e executable) writeWrapperArgv(wrapperPath string, destFolder string) error {
101+
if e.argv == nil {
102+
return nil
103+
}
94104
r := newReplacements(destDirPattern, destFolder)
95-
96-
// Add the shebang
97-
fmt.Fprintln(wrapper, "#! /bin/sh")
98-
99-
// Add the preceding lines if any
100-
for _, line := range e.preLines {
101-
fmt.Fprintf(wrapper, "%s\n", r.apply(line))
105+
f, err := os.OpenFile(wrapperPath+".argv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
106+
if err != nil {
107+
return err
102108
}
103-
104-
// Update the path to include the destination folder
105-
var env map[string]string
106-
if e.env == nil {
107-
env = make(map[string]string)
108-
} else {
109-
env = e.env
109+
defer f.Close()
110+
for _, arg := range e.argv {
111+
fmt.Fprintf(f, "%s\n", r.apply(arg))
110112
}
113+
return nil
114+
}
111115

112-
path, specified := env["PATH"]
113-
if !specified {
114-
path = "$PATH"
116+
func (e executable) writeWrapperEnvv(wrapperPath string, destFolder string) error {
117+
r := newReplacements(destDirPattern, destFolder)
118+
f, err := os.OpenFile(wrapperPath+".envv", os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0440)
119+
if err != nil {
120+
return err
115121
}
116-
env["PATH"] = strings.Join([]string{destFolder, path}, ":")
122+
defer f.Close()
117123

118-
var sortedEnvvars []string
119-
for e := range env {
120-
sortedEnvvars = append(sortedEnvvars, e)
124+
// Update PATH to insert the destination folder at the head.
125+
var envm map[string]string
126+
if e.envm == nil {
127+
envm = make(map[string]string)
128+
} else {
129+
envm = e.envm
121130
}
122-
sort.Strings(sortedEnvvars)
123-
124-
for _, e := range sortedEnvvars {
125-
v := env[e]
126-
fmt.Fprintf(wrapper, "%s=%s \\\n", e, r.apply(v))
131+
if path, ok := envm["PATH"]; ok {
132+
envm["PATH"] = destFolder + ":" + path
133+
} else {
134+
// Replace PATH with <PATH, which instructs wrapper to insert the value at the head of a
135+
// colon-separated environment variable list.
136+
delete(envm, "PATH")
137+
envm["<PATH"] = destFolder
127138
}
128-
// Add the call to the target executable
129-
fmt.Fprintf(wrapper, "%s \\\n", dotfileName)
130139

131-
// Insert additional lines in the `arg` list
132-
for _, line := range e.argLines {
133-
fmt.Fprintf(wrapper, "\t%s \\\n", r.apply(line))
140+
var envv []string
141+
for k, v := range envm {
142+
envv = append(envv, k+"="+r.apply(v))
143+
}
144+
sort.Strings(envv)
145+
for _, e := range envv {
146+
fmt.Fprintf(f, "%s\n", e)
134147
}
135-
// Add the script arguments "$@"
136-
fmt.Fprintln(wrapper, "\t\"$@\"")
137-
138148
return nil
139149
}
140150

‎tools/container/toolkit/executable_test.go

+83-56
Original file line numberDiff line numberDiff line change
@@ -17,102 +17,102 @@
1717
package main
1818

1919
import (
20-
"bytes"
20+
"bufio"
21+
"fmt"
22+
"io/fs"
2123
"os"
2224
"path/filepath"
23-
"strings"
2425
"testing"
2526

2627
"github.com/stretchr/testify/require"
2728
)
2829

2930
func TestWrapper(t *testing.T) {
30-
const shebang = "#! /bin/sh"
31-
const destFolder = "/dest/folder"
32-
const dotfileName = "source.real"
31+
createTestWrapperProgram(t)
3332

3433
testCases := []struct {
35-
e executable
36-
expectedLines []string
34+
e executable
35+
expectedArgv []string
36+
expectedEnvv []string
3737
}{
3838
{
39-
e: executable{},
40-
expectedLines: []string{
41-
shebang,
42-
"PATH=/dest/folder:$PATH \\",
43-
"source.real \\",
44-
"\t\"$@\"",
45-
"",
39+
e: executable{source: "source"},
40+
expectedEnvv: []string{
41+
fmt.Sprintf("<PATH=%s", destDirPattern),
4642
},
4743
},
4844
{
4945
e: executable{
50-
env: map[string]string{
51-
"PATH": "some-path",
46+
source: "source",
47+
envm: map[string]string{
48+
"FOO": "BAR",
5249
},
5350
},
54-
expectedLines: []string{
55-
shebang,
56-
"PATH=/dest/folder:some-path \\",
57-
"source.real \\",
58-
"\t\"$@\"",
59-
"",
51+
expectedEnvv: []string{
52+
fmt.Sprintf("<PATH=%s", destDirPattern),
53+
"FOO=BAR",
6054
},
6155
},
6256
{
6357
e: executable{
64-
preLines: []string{
65-
"preline1",
66-
"preline2",
58+
source: "source",
59+
envm: map[string]string{
60+
"PATH": "some-path",
61+
"FOO": "BAR",
6762
},
6863
},
69-
expectedLines: []string{
70-
shebang,
71-
"preline1",
72-
"preline2",
73-
"PATH=/dest/folder:$PATH \\",
74-
"source.real \\",
75-
"\t\"$@\"",
76-
"",
64+
expectedEnvv: []string{
65+
"FOO=BAR",
66+
fmt.Sprintf("PATH=%s:some-path", destDirPattern),
7767
},
7868
},
7969
{
8070
e: executable{
81-
argLines: []string{
82-
"argline1",
83-
"argline2",
71+
source: "source",
72+
argv: []string{
73+
"argb",
74+
"arga",
75+
"argc",
8476
},
8577
},
86-
expectedLines: []string{
87-
shebang,
88-
"PATH=/dest/folder:$PATH \\",
89-
"source.real \\",
90-
"\targline1 \\",
91-
"\targline2 \\",
92-
"\t\"$@\"",
93-
"",
78+
expectedArgv: []string{
79+
"argb",
80+
"arga",
81+
"argc",
82+
},
83+
expectedEnvv: []string{
84+
fmt.Sprintf("<PATH=%s", destDirPattern),
9485
},
9586
},
9687
}
9788

98-
for i, tc := range testCases {
99-
buf := &bytes.Buffer{}
100-
101-
err := tc.e.writeWrapperTo(buf, destFolder, dotfileName)
89+
for _, tc := range testCases {
90+
destFolder := t.TempDir()
91+
r := newReplacements(destDirPattern, destFolder)
92+
for k, v := range tc.expectedEnvv {
93+
tc.expectedEnvv[k] = r.apply(v)
94+
}
95+
path, err := tc.e.installWrapper(destFolder)
10296
require.NoError(t, err)
103-
104-
exepectedContents := strings.Join(tc.expectedLines, "\n")
105-
require.Equal(t, exepectedContents, buf.String(), "%v: %v", i, tc)
97+
require.FileExists(t, path)
98+
envv, err := readAllLines(path + ".envv")
99+
require.NoError(t, err)
100+
require.Equal(t, tc.expectedEnvv, envv)
101+
argv, err := readAllLines(path + ".argv")
102+
if tc.expectedArgv == nil {
103+
require.ErrorAs(t, err, &fs.ErrNotExist)
104+
} else {
105+
require.Equal(t, tc.expectedArgv, argv)
106+
107+
}
106108
}
107109
}
108110

109111
func TestInstallExecutable(t *testing.T) {
110-
inputFolder, err := os.MkdirTemp("", "")
111-
require.NoError(t, err)
112-
defer os.RemoveAll(inputFolder)
112+
createTestWrapperProgram(t)
113113

114114
// Create the source file
115-
source := filepath.Join(inputFolder, "input")
115+
source := filepath.Join(t.TempDir(), "input")
116116
sourceFile, err := os.Create(source)
117117

118118
base := filepath.Base(source)
@@ -123,7 +123,6 @@ func TestInstallExecutable(t *testing.T) {
123123
e := executable{
124124
source: source,
125125
target: executableTarget{
126-
dotfileName: "input.real",
127126
wrapperName: "input",
128127
},
129128
}
@@ -150,3 +149,31 @@ func TestInstallExecutable(t *testing.T) {
150149
require.NoError(t, err)
151150
require.NotEqual(t, 0, wrapperInfo.Mode()&0111)
152151
}
152+
153+
func createTestWrapperProgram(t *testing.T) {
154+
t.Helper()
155+
currentExe, err := os.Executable()
156+
if err != nil {
157+
t.Fatalf("error getting current executable: %v", err)
158+
}
159+
wrapperPath := filepath.Join(filepath.Dir(currentExe), "wrapper")
160+
f, err := os.Create(wrapperPath)
161+
if err != nil {
162+
t.Fatalf("error creating test wrapper: %v", err)
163+
}
164+
f.Close()
165+
}
166+
167+
func readAllLines(path string) (s []string, err error) {
168+
f, err := os.Open(path)
169+
if err != nil {
170+
return
171+
}
172+
defer f.Close()
173+
scanner := bufio.NewScanner(f)
174+
for scanner.Scan() {
175+
s = append(s, scanner.Text())
176+
}
177+
err = scanner.Err()
178+
return
179+
}

0 commit comments

Comments
 (0)