Skip to content

Commit f095bd5

Browse files
ArangoGutierrezArangoGutierrez
committed
Add --disable-hook flag to cdi generate command
When running the nvidia-ctk cdi generate command, a user should be able to opt out of specific hooks. We propose to add a flag --disable-hook that will take a comma-separated list of hooks that will be skipped when creating the CDI spec. Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>
1 parent aa696ef commit f095bd5

File tree

13 files changed

+225
-45
lines changed

13 files changed

+225
-45
lines changed

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

Lines changed: 16 additions & 2 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
@@ -176,6 +177,13 @@ func (m command) build() *cli.Command {
176177
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: "Hook to skip when generating the CDI specification. Can be specified multiple times. Can be a comma-separated list of hooks or a single hook name.",
184+
Value: cli.NewStringSlice(),
185+
Destination: &opts.disabledHooks,
186+
},
179187
}
180188

181189
return &c
@@ -262,7 +270,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
262270
deviceNamers = append(deviceNamers, deviceNamer)
263271
}
264272

265-
cdilib, err := nvcdi.New(
273+
initOpts := []nvcdi.Option{
266274
nvcdi.WithLogger(m.logger),
267275
nvcdi.WithDriverRoot(opts.driverRoot),
268276
nvcdi.WithDevRoot(opts.devRoot),
@@ -276,7 +284,13 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
276284
nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
277285
// We set the following to allow for dependency injection:
278286
nvcdi.WithNvmlLib(opts.nvmllib),
279-
)
287+
}
288+
289+
for _, hook := range opts.disabledHooks.Value() {
290+
initOpts = append(initOpts, nvcdi.WithDisabledHook(hook))
291+
}
292+
293+
cdilib, err := nvcdi.New(initOpts...)
280294
if err != nil {
281295
return nil, fmt.Errorf("failed to create CDI library: %v", err)
282296
}

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

Lines changed: 177 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
)
@@ -36,6 +37,9 @@ func TestGenerateSpec(t *testing.T) {
3637
require.NoError(t, err)
3738

3839
driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")
40+
disableHook1 := cli.NewStringSlice("enable-cuda-compat")
41+
disableHook2 := cli.NewStringSlice("enable-cuda-compat", "update-ldcache")
42+
disableHook3 := cli.NewStringSlice("all")
3943

4044
logger, _ := testlog.NewNullLogger()
4145
testCases := []struct {
@@ -113,6 +117,179 @@ containerEdits:
113117
- nodev
114118
- rbind
115119
- rprivate
120+
`,
121+
},
122+
{
123+
description: "disableHooks1",
124+
options: options{
125+
format: "yaml",
126+
mode: "nvml",
127+
vendor: "example.com",
128+
class: "device",
129+
driverRoot: driverRoot,
130+
disabledHooks: *disableHook1,
131+
},
132+
expectedOptions: options{
133+
format: "yaml",
134+
mode: "nvml",
135+
vendor: "example.com",
136+
class: "device",
137+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
138+
driverRoot: driverRoot,
139+
disabledHooks: *disableHook1,
140+
},
141+
expectedSpec: `---
142+
cdiVersion: 0.5.0
143+
kind: example.com/device
144+
devices:
145+
- name: "0"
146+
containerEdits:
147+
deviceNodes:
148+
- path: /dev/nvidia0
149+
hostPath: {{ .driverRoot }}/dev/nvidia0
150+
- name: all
151+
containerEdits:
152+
deviceNodes:
153+
- path: /dev/nvidia0
154+
hostPath: {{ .driverRoot }}/dev/nvidia0
155+
containerEdits:
156+
env:
157+
- NVIDIA_VISIBLE_DEVICES=void
158+
deviceNodes:
159+
- path: /dev/nvidiactl
160+
hostPath: {{ .driverRoot }}/dev/nvidiactl
161+
hooks:
162+
- hookName: createContainer
163+
path: /usr/bin/nvidia-cdi-hook
164+
args:
165+
- nvidia-cdi-hook
166+
- create-symlinks
167+
- --link
168+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
169+
- hookName: createContainer
170+
path: /usr/bin/nvidia-cdi-hook
171+
args:
172+
- nvidia-cdi-hook
173+
- update-ldcache
174+
- --folder
175+
- /lib/x86_64-linux-gnu
176+
mounts:
177+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
178+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
179+
options:
180+
- ro
181+
- nosuid
182+
- nodev
183+
- rbind
184+
- rprivate
185+
`,
186+
},
187+
{
188+
description: "disableHooks2",
189+
options: options{
190+
format: "yaml",
191+
mode: "nvml",
192+
vendor: "example.com",
193+
class: "device",
194+
driverRoot: driverRoot,
195+
disabledHooks: *disableHook2,
196+
},
197+
expectedOptions: options{
198+
format: "yaml",
199+
mode: "nvml",
200+
vendor: "example.com",
201+
class: "device",
202+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
203+
driverRoot: driverRoot,
204+
disabledHooks: *disableHook2,
205+
},
206+
expectedSpec: `---
207+
cdiVersion: 0.5.0
208+
kind: example.com/device
209+
devices:
210+
- name: "0"
211+
containerEdits:
212+
deviceNodes:
213+
- path: /dev/nvidia0
214+
hostPath: {{ .driverRoot }}/dev/nvidia0
215+
- name: all
216+
containerEdits:
217+
deviceNodes:
218+
- path: /dev/nvidia0
219+
hostPath: {{ .driverRoot }}/dev/nvidia0
220+
containerEdits:
221+
env:
222+
- NVIDIA_VISIBLE_DEVICES=void
223+
deviceNodes:
224+
- path: /dev/nvidiactl
225+
hostPath: {{ .driverRoot }}/dev/nvidiactl
226+
hooks:
227+
- hookName: createContainer
228+
path: /usr/bin/nvidia-cdi-hook
229+
args:
230+
- nvidia-cdi-hook
231+
- create-symlinks
232+
- --link
233+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
234+
mounts:
235+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
236+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
237+
options:
238+
- ro
239+
- nosuid
240+
- nodev
241+
- rbind
242+
- rprivate
243+
`,
244+
},
245+
{
246+
description: "disableHooksAll",
247+
options: options{
248+
format: "yaml",
249+
mode: "nvml",
250+
vendor: "example.com",
251+
class: "device",
252+
driverRoot: driverRoot,
253+
disabledHooks: *disableHook3,
254+
},
255+
expectedOptions: options{
256+
format: "yaml",
257+
mode: "nvml",
258+
vendor: "example.com",
259+
class: "device",
260+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
261+
driverRoot: driverRoot,
262+
disabledHooks: *disableHook3,
263+
},
264+
expectedSpec: `---
265+
cdiVersion: 0.5.0
266+
kind: example.com/device
267+
devices:
268+
- name: "0"
269+
containerEdits:
270+
deviceNodes:
271+
- path: /dev/nvidia0
272+
hostPath: {{ .driverRoot }}/dev/nvidia0
273+
- name: all
274+
containerEdits:
275+
deviceNodes:
276+
- path: /dev/nvidia0
277+
hostPath: {{ .driverRoot }}/dev/nvidia0
278+
containerEdits:
279+
env:
280+
- NVIDIA_VISIBLE_DEVICES=void
281+
deviceNodes:
282+
- path: /dev/nvidiactl
283+
hostPath: {{ .driverRoot }}/dev/nvidiactl
284+
mounts:
285+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
286+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
287+
options:
288+
- ro
289+
- nosuid
290+
- nodev
291+
- rbind
292+
- rprivate
116293
`,
117294
},
118295
}

‎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")
28+
hookCreator := NewHookCreator("/usr/bin/nvidia-cdi-hook", make(DisabledHooks))
2929

3030
testCases := []struct {
3131
description string

‎internal/discover/hooks.go

Lines changed: 18 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -49,6 +49,17 @@ type HookName string
4949
// DisabledHooks allows individual hooks to be disabled.
5050
type DisabledHooks map[HookName]bool
5151

52+
func (d DisabledHooks) Set(value HookName) {
53+
if value == "all" {
54+
for _, hook := range AllHooks {
55+
d[hook] = true
56+
}
57+
return
58+
}
59+
60+
d[value] = true
61+
}
62+
5263
const (
5364
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
5465
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
@@ -72,21 +83,27 @@ type Option func(*CDIHook)
7283

7384
type CDIHook struct {
7485
nvidiaCDIHookPath string
86+
disabledHooks DisabledHooks
7587
}
7688

7789
type HookCreator interface {
7890
Create(HookName, ...string) *Hook
7991
}
8092

81-
func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
93+
func NewHookCreator(nvidiaCDIHookPath string, disabledHooks DisabledHooks) HookCreator {
8294
CDIHook := &CDIHook{
8395
nvidiaCDIHookPath: nvidiaCDIHookPath,
96+
disabledHooks: disabledHooks,
8497
}
8598

8699
return CDIHook
87100
}
88101

89102
func (c CDIHook) Create(name HookName, args ...string) *Hook {
103+
if c.disabledHooks[name] {
104+
return nil
105+
}
106+
90107
if name == "create-symlinks" {
91108
if len(args) == 0 {
92109
return nil

‎internal/discover/ldconfig_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ const (
3131

3232
func TestLDCacheUpdateHook(t *testing.T) {
3333
logger, _ := testlog.NewNullLogger()
34-
hookCreator := NewHookCreator(testNvidiaCDIHookPath)
34+
hookCreator := NewHookCreator(testNvidiaCDIHookPath, make(DisabledHooks))
3535

3636
testCases := []struct {
3737
description string

‎internal/discover/symlinks_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -306,7 +306,7 @@ func TestWithWithDriverDotSoSymlinks(t *testing.T) {
306306
},
307307
}
308308

309-
hookCreator := NewHookCreator("/path/to/nvidia-cdi-hook")
309+
hookCreator := NewHookCreator("/path/to/nvidia-cdi-hook", make(DisabledHooks))
310310
for _, tc := range testCases {
311311
t.Run(tc.description, func(t *testing.T) {
312312
d := WithDriverDotSoSymlinks(

‎internal/platform-support/tegra/csv_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,7 +181,7 @@ func TestDiscovererFromCSVFiles(t *testing.T) {
181181
},
182182
}
183183

184-
hookCreator := discover.NewHookCreator("/usr/bin/nvidia-cdi-hook")
184+
hookCreator := discover.NewHookCreator("/usr/bin/nvidia-cdi-hook", make(discover.DisabledHooks))
185185
for _, tc := range testCases {
186186
t.Run(tc.description, func(t *testing.T) {
187187
defer setGetTargetsFromCSVFiles(tc.moutSpecs)()

‎internal/runtime/runtime_factory.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ func newSpecModifier(logger logger.Interface, cfg *config.Config, ociSpec oci.Sp
7575
return nil, err
7676
}
7777

78-
hookCreator := discover.NewHookCreator(cfg.NVIDIACTKConfig.Path)
78+
hookCreator := discover.NewHookCreator(cfg.NVIDIACTKConfig.Path, make(discover.DisabledHooks))
7979

8080
mode := info.ResolveAutoMode(logger, cfg.NVIDIAContainerRuntimeConfig.Mode, image)
8181
// We update the mode here so that we can continue passing just the config to other functions.

‎pkg/nvcdi/driver-nvml.go

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -106,11 +106,9 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover
106106
)
107107
discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)
108108

109-
if l.HookIsSupported(HookEnableCudaCompat) {
110-
// TODO: The following should use the version directly.
111-
cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, l.driver)
112-
discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
113-
}
109+
// TODO: The following should use the version directly.
110+
cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, l.driver)
111+
discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
114112

115113
updateLDCache, _ := discover.NewLDCacheUpdateHook(l.logger, libraries, l.hookCreator, l.ldconfigPath)
116114
discoverers = append(discoverers, updateLDCache)

‎pkg/nvcdi/driver-wsl_test.go

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

3030
func TestNvidiaSMISymlinkHook(t *testing.T) {
3131
logger, _ := testlog.NewNullLogger()
32-
hookCreator := discover.NewHookCreator("nvidia-cdi-hook")
32+
hookCreator := discover.NewHookCreator("nvidia-cdi-hook", make(discover.DisabledHooks))
3333

3434
errMounts := errors.New("mounts error")
3535

‎pkg/nvcdi/hooks.go

Lines changed: 0 additions & 27 deletions
This file was deleted.

0 commit comments

Comments
 (0)