Skip to content

Commit 191f11e

Browse files
ArangoGutierrezArangoGutierrez
committed
Add all option to disable all hooks
Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>
1 parent d657e0a commit 191f11e

File tree

6 files changed

+166
-20
lines changed

6 files changed

+166
-20
lines changed

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

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

5858
configSearchPaths cli.StringSlice
5959
librarySearchPaths cli.StringSlice
60-
disableHooks cli.StringSlice
60+
disabledHooks cli.StringSlice
6161

6262
csv struct {
6363
files cli.StringSlice
@@ -179,9 +179,10 @@ func (m command) build() *cli.Command {
179179
},
180180
&cli.StringSliceFlag{
181181
Name: "disable-hook",
182-
Usage: "Comma-separated list of hooks to skip when generating the CDI specification.",
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.",
183184
Value: cli.NewStringSlice(),
184-
Destination: &opts.disableHooks,
185+
Destination: &opts.disabledHooks,
185186
},
186187
}
187188

@@ -285,10 +286,8 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
285286
nvcdi.WithNvmlLib(opts.nvmllib),
286287
}
287288

288-
if len(opts.disableHooks.Value()) > 0 {
289-
for _, hook := range opts.disableHooks.Value() {
290-
initOpts = append(initOpts, nvcdi.WithDisabledHook(nvcdi.HookName(hook)))
291-
}
289+
for _, hook := range opts.disabledHooks.Value() {
290+
initOpts = append(initOpts, nvcdi.WithDisabledHook(hook))
292291
}
293292

294293
cdilib, err := nvcdi.New(initOpts...)

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

Lines changed: 132 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,9 @@ func TestGenerateSpec(t *testing.T) {
3737
require.NoError(t, err)
3838

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

4244
logger, _ := testlog.NewNullLogger()
4345
testCases := []struct {
@@ -117,14 +119,14 @@ containerEdits:
117119
`,
118120
},
119121
{
120-
description: "disableHooks",
122+
description: "disableHooks1",
121123
options: options{
122-
format: "yaml",
123-
mode: "nvml",
124-
vendor: "example.com",
125-
class: "device",
126-
driverRoot: driverRoot,
127-
disableHooks: *disableHook,
124+
format: "yaml",
125+
mode: "nvml",
126+
vendor: "example.com",
127+
class: "device",
128+
driverRoot: driverRoot,
129+
disabledHooks: *disableHook1,
128130
},
129131
expectedOptions: options{
130132
format: "yaml",
@@ -133,7 +135,128 @@ containerEdits:
133135
class: "device",
134136
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
135137
driverRoot: driverRoot,
136-
disableHooks: *disableHook,
138+
disabledHooks: *disableHook1,
139+
},
140+
expectedSpec: `---
141+
cdiVersion: 0.5.0
142+
kind: example.com/device
143+
devices:
144+
- name: "0"
145+
containerEdits:
146+
deviceNodes:
147+
- path: /dev/nvidia0
148+
hostPath: {{ .driverRoot }}/dev/nvidia0
149+
- name: all
150+
containerEdits:
151+
deviceNodes:
152+
- path: /dev/nvidia0
153+
hostPath: {{ .driverRoot }}/dev/nvidia0
154+
containerEdits:
155+
env:
156+
- NVIDIA_VISIBLE_DEVICES=void
157+
deviceNodes:
158+
- path: /dev/nvidiactl
159+
hostPath: {{ .driverRoot }}/dev/nvidiactl
160+
hooks:
161+
- hookName: createContainer
162+
path: /usr/bin/nvidia-cdi-hook
163+
args:
164+
- nvidia-cdi-hook
165+
- create-symlinks
166+
- --link
167+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
168+
- hookName: createContainer
169+
path: /usr/bin/nvidia-cdi-hook
170+
args:
171+
- nvidia-cdi-hook
172+
- update-ldcache
173+
- --folder
174+
- /lib/x86_64-linux-gnu
175+
mounts:
176+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
177+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
178+
options:
179+
- ro
180+
- nosuid
181+
- nodev
182+
- bind
183+
`,
184+
},
185+
{
186+
description: "disableHooks2",
187+
options: options{
188+
format: "yaml",
189+
mode: "nvml",
190+
vendor: "example.com",
191+
class: "device",
192+
driverRoot: driverRoot,
193+
disabledHooks: *disableHook2,
194+
},
195+
expectedOptions: options{
196+
format: "yaml",
197+
mode: "nvml",
198+
vendor: "example.com",
199+
class: "device",
200+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
201+
driverRoot: driverRoot,
202+
disabledHooks: *disableHook2,
203+
},
204+
expectedSpec: `---
205+
cdiVersion: 0.5.0
206+
kind: example.com/device
207+
devices:
208+
- name: "0"
209+
containerEdits:
210+
deviceNodes:
211+
- path: /dev/nvidia0
212+
hostPath: {{ .driverRoot }}/dev/nvidia0
213+
- name: all
214+
containerEdits:
215+
deviceNodes:
216+
- path: /dev/nvidia0
217+
hostPath: {{ .driverRoot }}/dev/nvidia0
218+
containerEdits:
219+
env:
220+
- NVIDIA_VISIBLE_DEVICES=void
221+
deviceNodes:
222+
- path: /dev/nvidiactl
223+
hostPath: {{ .driverRoot }}/dev/nvidiactl
224+
hooks:
225+
- hookName: createContainer
226+
path: /usr/bin/nvidia-cdi-hook
227+
args:
228+
- nvidia-cdi-hook
229+
- create-symlinks
230+
- --link
231+
- libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
232+
mounts:
233+
- hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
234+
containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
235+
options:
236+
- ro
237+
- nosuid
238+
- nodev
239+
- bind
240+
`,
241+
},
242+
{
243+
description: "disableHooksAll",
244+
options: options{
245+
format: "yaml",
246+
mode: "nvml",
247+
vendor: "example.com",
248+
class: "device",
249+
driverRoot: driverRoot,
250+
disabledHooks: *disableHook3,
251+
},
252+
expectedOptions: options{
253+
format: "yaml",
254+
mode: "nvml",
255+
vendor: "example.com",
256+
class: "device",
257+
nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
258+
driverRoot: driverRoot,
259+
disabledHooks: *disableHook3,
137260
},
138261
expectedSpec: `---
139262
cdiVersion: 0.5.0

‎pkg/nvcdi/api.go

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -44,4 +44,17 @@ const (
4444
// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
4545
// This was added with v1.17.5 of the NVIDIA Container Toolkit.
4646
HookEnableCudaCompat = HookName("enable-cuda-compat")
47+
// HookCreateSymlinks refers to the hook used create symlinks inside the
48+
// directory path to be mounted into a container.
49+
HookCreateSymlinks = HookName("create-symlinks")
50+
// HookUpdateLDCache refers to the hook used to Update the dynamic linker
51+
// cache inside the directory path to be mounted into a container.
52+
HookUpdateLDCache = HookName("update-ldcache")
4753
)
54+
55+
// AllHooks maintains a future-proof list of all defined hooks.
56+
var AllHooks = []HookName{
57+
HookEnableCudaCompat,
58+
HookCreateSymlinks,
59+
HookUpdateLDCache,
60+
}

‎pkg/nvcdi/driver-nvml.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,7 +99,7 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover
9999

100100
var discoverers []discover.Discover
101101

102-
if l.HookIsSupported("create-symlinks") {
102+
if l.HookIsSupported(HookCreateSymlinks) {
103103
driverDotSoSymlinksDiscoverer := discover.WithDriverDotSoSymlinks(
104104
libraries,
105105
version,
@@ -118,7 +118,7 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover
118118
discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
119119
}
120120

121-
if l.HookIsSupported("update-ldcache") {
121+
if l.HookIsSupported(HookUpdateLDCache) {
122122
updateLDCache, _ := discover.NewLDCacheUpdateHook(
123123
l.logger,
124124
libraries,

‎pkg/nvcdi/hooks.go

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,3 +28,14 @@ func (l *nvcdilib) HookIsSupported(h HookName) bool {
2828
}
2929
return !l.disabledHooks[h]
3030
}
31+
32+
func (d disabledHooks) Set(value HookName) {
33+
if value == "all" {
34+
for _, hook := range AllHooks {
35+
d[hook] = true
36+
}
37+
return
38+
}
39+
40+
d[value] = true
41+
}

‎pkg/nvcdi/options.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -158,11 +158,11 @@ func WithLibrarySearchPaths(paths []string) Option {
158158

159159
// WithDisabledHook allows specific hooks to the disabled.
160160
// This option can be specified multiple times for each hook.
161-
func WithDisabledHook(hook HookName) Option {
161+
func WithDisabledHook[T string | HookName](hook T) Option {
162162
return func(o *nvcdilib) {
163163
if o.disabledHooks == nil {
164164
o.disabledHooks = make(map[HookName]bool)
165165
}
166-
o.disabledHooks[hook] = true
166+
o.disabledHooks.Set(HookName(hook))
167167
}
168168
}

0 commit comments

Comments
 (0)