Skip to content

Commit 6cf0248

Browse files
ArangoGutierrezelezar
authored andcommitted
Added ability to disable specific (or all) CDI hooks
This change adds the ability to disabled specific (or all) CDI hooks to both the nvidia-ctk cdi generate command and the nvcdi API. Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]> Signed-off-by: Evan Lezar <[email protected]>
1 parent b478751 commit 6cf0248

File tree

16 files changed

+372
-119
lines changed

16 files changed

+372
-119
lines changed

cmd/nvidia-ctk/cdi/generate/generate.go

Lines changed: 19 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@ type options struct {
5757

5858
configSearchPaths cli.StringSlice
5959
librarySearchPaths cli.StringSlice
60+
disabledHooks cli.StringSlice
6061

6162
csv struct {
6263
files cli.StringSlice
@@ -173,9 +174,18 @@ func (m command) build() *cli.Command {
173174
},
174175
&cli.StringSliceFlag{
175176
Name: "csv.ignore-pattern",
176-
Usage: "Specify a pattern the CSV mount specifications.",
177+
Usage: "specify a pattern the CSV mount specifications.",
177178
Destination: &opts.csv.ignorePatterns,
178179
},
180+
&cli.StringSliceFlag{
181+
Name: "disable-hook",
182+
Aliases: []string{"disable-hooks"},
183+
Usage: "specify a specific hook to skip when generating CDI " +
184+
"specifications. This can be specified multiple times and the " +
185+
"special hook name 'all' can be used ensure that the generated " +
186+
"CDI specification does not include any hooks.",
187+
Destination: &opts.disabledHooks,
188+
},
179189
}
180190

181191
return &c
@@ -262,7 +272,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
262272
deviceNamers = append(deviceNamers, deviceNamer)
263273
}
264274

265-
cdilib, err := nvcdi.New(
275+
cdiOptions := []nvcdi.Option{
266276
nvcdi.WithLogger(m.logger),
267277
nvcdi.WithDriverRoot(opts.driverRoot),
268278
nvcdi.WithDevRoot(opts.devRoot),
@@ -276,7 +286,13 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
276286
nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
277287
// We set the following to allow for dependency injection:
278288
nvcdi.WithNvmlLib(opts.nvmllib),
279-
)
289+
}
290+
291+
for _, hook := range opts.disabledHooks.Value() {
292+
cdiOptions = append(cdiOptions, nvcdi.WithDisabledHook(hook))
293+
}
294+
295+
cdilib, err := nvcdi.New(cdiOptions...)
280296
if err != nil {
281297
return nil, fmt.Errorf("failed to create CDI library: %v", err)
282298
}

cmd/nvidia-ctk/cdi/generate/generate_test.go

Lines changed: 186 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@ import (
2626
"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
2727
testlog "github.com/sirupsen/logrus/hooks/test"
2828
"github.com/stretchr/testify/require"
29+
"github.com/urfave/cli/v2"
2930

3031
"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
3132
)
@@ -119,6 +120,185 @@ containerEdits:
119120
- nodev
120121
- rbind
121122
- rprivate
123+
`,
124+
},
125+
{
126+
description: "disableHooks1",
127+
options: options{
128+
format: "yaml",
129+
mode: "nvml",
130+
vendor: "example.com",
131+
class: "device",
132+
driverRoot: driverRoot,
133+
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat")),
134+
},
135+
expectedOptions: options{
136+
format: "yaml",
137+
mode: "nvml",
138+
vendor: "example.com",
139+
class: "device",
140+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
141+
driverRoot: driverRoot,
142+
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat")),
143+
},
144+
expectedSpec: `---
145+
cdiVersion: 0.5.0
146+
kind: example.com/device
147+
devices:
148+
- name: "0"
149+
containerEdits:
150+
deviceNodes:
151+
- path: /dev/nvidia0
152+
hostPath: {{ .driverRoot }}/dev/nvidia0
153+
- name: all
154+
containerEdits:
155+
deviceNodes:
156+
- path: /dev/nvidia0
157+
hostPath: {{ .driverRoot }}/dev/nvidia0
158+
containerEdits:
159+
env:
160+
- NVIDIA_VISIBLE_DEVICES=void
161+
deviceNodes:
162+
- path: /dev/nvidiactl
163+
hostPath: {{ .driverRoot }}/dev/nvidiactl
164+
hooks:
165+
- hookName: createContainer
166+
path: /usr/bin/nvidia-cdi-hook
167+
args:
168+
- nvidia-cdi-hook
169+
- create-symlinks
170+
- --link
171+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
172+
env:
173+
- NVIDIA_CTK_DEBUG=false
174+
- hookName: createContainer
175+
path: /usr/bin/nvidia-cdi-hook
176+
args:
177+
- nvidia-cdi-hook
178+
- update-ldcache
179+
- --folder
180+
- /lib/x86_64-linux-gnu
181+
env:
182+
- NVIDIA_CTK_DEBUG=false
183+
mounts:
184+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
185+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
186+
options:
187+
- ro
188+
- nosuid
189+
- nodev
190+
- rbind
191+
- rprivate
192+
`,
193+
},
194+
{
195+
description: "disableHooks2",
196+
options: options{
197+
format: "yaml",
198+
mode: "nvml",
199+
vendor: "example.com",
200+
class: "device",
201+
driverRoot: driverRoot,
202+
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat", "update-ldcache")),
203+
},
204+
expectedOptions: options{
205+
format: "yaml",
206+
mode: "nvml",
207+
vendor: "example.com",
208+
class: "device",
209+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
210+
driverRoot: driverRoot,
211+
disabledHooks: valueOf(cli.NewStringSlice("enable-cuda-compat", "update-ldcache")),
212+
},
213+
expectedSpec: `---
214+
cdiVersion: 0.5.0
215+
kind: example.com/device
216+
devices:
217+
- name: "0"
218+
containerEdits:
219+
deviceNodes:
220+
- path: /dev/nvidia0
221+
hostPath: {{ .driverRoot }}/dev/nvidia0
222+
- name: all
223+
containerEdits:
224+
deviceNodes:
225+
- path: /dev/nvidia0
226+
hostPath: {{ .driverRoot }}/dev/nvidia0
227+
containerEdits:
228+
env:
229+
- NVIDIA_VISIBLE_DEVICES=void
230+
deviceNodes:
231+
- path: /dev/nvidiactl
232+
hostPath: {{ .driverRoot }}/dev/nvidiactl
233+
hooks:
234+
- hookName: createContainer
235+
path: /usr/bin/nvidia-cdi-hook
236+
args:
237+
- nvidia-cdi-hook
238+
- create-symlinks
239+
- --link
240+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
241+
env:
242+
- NVIDIA_CTK_DEBUG=false
243+
mounts:
244+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
245+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
246+
options:
247+
- ro
248+
- nosuid
249+
- nodev
250+
- rbind
251+
- rprivate
252+
`,
253+
},
254+
{
255+
description: "disableHooksAll",
256+
options: options{
257+
format: "yaml",
258+
mode: "nvml",
259+
vendor: "example.com",
260+
class: "device",
261+
driverRoot: driverRoot,
262+
disabledHooks: valueOf(cli.NewStringSlice("all")),
263+
},
264+
expectedOptions: options{
265+
format: "yaml",
266+
mode: "nvml",
267+
vendor: "example.com",
268+
class: "device",
269+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
270+
driverRoot: driverRoot,
271+
disabledHooks: valueOf(cli.NewStringSlice("all")),
272+
},
273+
expectedSpec: `---
274+
cdiVersion: 0.5.0
275+
kind: example.com/device
276+
devices:
277+
- name: "0"
278+
containerEdits:
279+
deviceNodes:
280+
- path: /dev/nvidia0
281+
hostPath: {{ .driverRoot }}/dev/nvidia0
282+
- name: all
283+
containerEdits:
284+
deviceNodes:
285+
- path: /dev/nvidia0
286+
hostPath: {{ .driverRoot }}/dev/nvidia0
287+
containerEdits:
288+
env:
289+
- NVIDIA_VISIBLE_DEVICES=void
290+
deviceNodes:
291+
- path: /dev/nvidiactl
292+
hostPath: {{ .driverRoot }}/dev/nvidiactl
293+
mounts:
294+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
295+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
296+
options:
297+
- ro
298+
- nosuid
299+
- nodev
300+
- rbind
301+
- rprivate
122302
`,
123303
},
124304
}
@@ -162,3 +342,9 @@ containerEdits:
162342
})
163343
}
164344
}
345+
346+
// valueOf returns the value of a pointer.
347+
// Note that this does not check for a nil pointer and is only used for testing.
348+
func valueOf[T any](v *T) T {
349+
return *v
350+
}

internal/discover/graphics_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ import (
2525

2626
func TestGraphicsLibrariesDiscoverer(t *testing.T) {
2727
logger, _ := testlog.NewNullLogger()
28-
hookCreator := NewHookCreator("/usr/bin/nvidia-cdi-hook", false)
28+
hookCreator := NewHookCreator()
2929

3030
testCases := []struct {
3131
description string

0 commit comments

Comments
 (0)