1818package installer
1919
2020import (
21- "bytes"
22- "fmt"
23- "html/template"
24- "io"
2521 "path/filepath"
26- "strings"
2722
2823 log "github.com/sirupsen/logrus"
2924
3025 "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/container/operator"
26+ "github.com/NVIDIA/nvidia-container-toolkit/cmd/nvidia-ctk-installer/toolkit/installer/wrappercore"
3127 "github.com/NVIDIA/nvidia-container-toolkit/internal/config"
3228)
3329
3430type executable struct {
3531 requiresKernelModule bool
3632 path string
3733 symlink string
38- env map [string ]string
34+ argv []string
35+ envm map [string ]string
3936}
4037
4138func (t * ToolkitInstaller ) collectExecutables (destDir string ) ([]Installer , error ) {
@@ -53,7 +50,7 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
5350 e := executable {
5451 path : runtime .Path ,
5552 requiresKernelModule : true ,
56- env : map [string ]string {
53+ envm : map [string ]string {
5754 config .FilePathOverrideEnvVar : configFilePath ,
5855 },
5956 }
@@ -62,15 +59,15 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
6259 executables = append (executables ,
6360 executable {
6461 path : "nvidia-container-cli" ,
65- env : map [string ]string {"LD_LIBRARY_PATH" : destDir + ":$LD_LIBRARY_PATH" },
62+ envm : map [string ]string {"< LD_LIBRARY_PATH" : destDir },
6663 },
6764 )
6865
6966 executables = append (executables ,
7067 executable {
7168 path : "nvidia-container-runtime-hook" ,
7269 symlink : "nvidia-container-toolkit" ,
73- env : map [string ]string {
70+ envm : map [string ]string {
7471 config .FilePathOverrideEnvVar : configFilePath ,
7572 },
7673 },
@@ -87,25 +84,25 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
8784 return nil , err
8885 }
8986
90- wrappedExecutableFilename := filepath .Base (executablePath )
91- dotRealFilename := wrappedExecutableFilename + ".real"
92-
9387 w := & wrapper {
94- Source : executablePath ,
95- WrappedExecutable : dotRealFilename ,
96- CheckModules : executable .requiresKernelModule ,
97- Envvars : map [string ]string {
98- "PATH" : strings .Join ([]string {destDir , "$PATH" }, ":" ),
88+ Source : executablePath ,
89+ WrapperProgramPath : t .wrapperProgramPath ,
90+ Config : wrappercore.WrapperConfig {
91+ Argv : executable .argv ,
92+ Envm : map [string ]string {
93+ "<PATH" : destDir ,
94+ },
95+ RequiresKernelModule : executable .requiresKernelModule ,
9996 },
10097 }
101- for k , v := range executable .env {
102- w .Envvars [k ] = v
98+ for k , v := range executable .envm {
99+ w .Config . Envm [k ] = v
103100 }
104101
105102 if len (t .defaultRuntimeExecutablePath ) > 0 {
106- w .DefaultRuntimeExecutablePath = t .defaultRuntimeExecutablePath
103+ w .Config . DefaultRuntimeExecutablePath = t .defaultRuntimeExecutablePath
107104 } else {
108- w .DefaultRuntimeExecutablePath = "runc"
105+ w .Config . DefaultRuntimeExecutablePath = "runc"
109106 }
110107
111108 installers = append (installers , w )
@@ -121,66 +118,4 @@ func (t *ToolkitInstaller) collectExecutables(destDir string) ([]Installer, erro
121118 }
122119
123120 return installers , nil
124-
125- }
126-
127- type wrapper struct {
128- Source string
129- Envvars map [string ]string
130- WrappedExecutable string
131- CheckModules bool
132- DefaultRuntimeExecutablePath string
133- }
134-
135- type render struct {
136- * wrapper
137- DestDir string
138- }
139-
140- func (w * wrapper ) Install (destDir string ) error {
141- // Copy the executable with a .real extension.
142- mode , err := installFile (w .Source , filepath .Join (destDir , w .WrappedExecutable ))
143- if err != nil {
144- return err
145- }
146-
147- // Create a wrapper file.
148- r := render {
149- wrapper : w ,
150- DestDir : destDir ,
151- }
152- content , err := r .render ()
153- if err != nil {
154- return fmt .Errorf ("failed to render wrapper: %w" , err )
155- }
156- wrapperFile := filepath .Join (destDir , filepath .Base (w .Source ))
157- return installContent (content , wrapperFile , mode | 0111 )
158- }
159-
160- func (w * render ) render () (io.Reader , error ) {
161- wrapperTemplate := `#! /bin/sh
162- {{- if (.CheckModules) }}
163- cat /proc/modules | grep -e "^nvidia " >/dev/null 2>&1
164- if [ "${?}" != "0" ]; then
165- echo "nvidia driver modules are not yet loaded, invoking {{ .DefaultRuntimeExecutablePath }} directly" >&2
166- exec {{ .DefaultRuntimeExecutablePath }} "$@"
167- fi
168- {{- end }}
169- {{- range $key, $value := .Envvars }}
170- {{$key}}={{$value}} \
171- {{- end }}
172- {{ .DestDir }}/{{ .WrappedExecutable }} \
173- "$@"
174- `
175-
176- var content bytes.Buffer
177- tmpl , err := template .New ("wrapper" ).Parse (wrapperTemplate )
178- if err != nil {
179- return nil , err
180- }
181- if err := tmpl .Execute (& content , w ); err != nil {
182- return nil , err
183- }
184-
185- return & content , nil
186121}
0 commit comments