From 191f11e3f9d934c607ddc4bfc47438a74bff9823 Mon Sep 17 00:00:00 2001 From: Carlos Eduardo Arango Gutierrez Date: Wed, 14 May 2025 12:26:00 +0200 Subject: [PATCH] Add all option to disable all hooks Signed-off-by: Carlos Eduardo Arango Gutierrez --- cmd/nvidia-ctk/cdi/generate/generate.go | 13 +- cmd/nvidia-ctk/cdi/generate/generate_test.go | 141 +++++++++++++++++-- pkg/nvcdi/api.go | 13 ++ pkg/nvcdi/driver-nvml.go | 4 +- pkg/nvcdi/hooks.go | 11 ++ pkg/nvcdi/options.go | 4 +- 6 files changed, 166 insertions(+), 20 deletions(-) diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go index 14fa6af0..569cd133 100644 --- a/cmd/nvidia-ctk/cdi/generate/generate.go +++ b/cmd/nvidia-ctk/cdi/generate/generate.go @@ -57,7 +57,7 @@ type options struct { configSearchPaths cli.StringSlice librarySearchPaths cli.StringSlice - disableHooks cli.StringSlice + disabledHooks cli.StringSlice csv struct { files cli.StringSlice @@ -179,9 +179,10 @@ func (m command) build() *cli.Command { }, &cli.StringSliceFlag{ Name: "disable-hook", - Usage: "Comma-separated list of hooks to skip when generating the CDI specification.", + Aliases: []string{"disable-hooks"}, + 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.", Value: cli.NewStringSlice(), - Destination: &opts.disableHooks, + Destination: &opts.disabledHooks, }, } @@ -285,10 +286,8 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) { nvcdi.WithNvmlLib(opts.nvmllib), } - if len(opts.disableHooks.Value()) > 0 { - for _, hook := range opts.disableHooks.Value() { - initOpts = append(initOpts, nvcdi.WithDisabledHook(nvcdi.HookName(hook))) - } + for _, hook := range opts.disabledHooks.Value() { + initOpts = append(initOpts, nvcdi.WithDisabledHook(hook)) } cdilib, err := nvcdi.New(initOpts...) diff --git a/cmd/nvidia-ctk/cdi/generate/generate_test.go b/cmd/nvidia-ctk/cdi/generate/generate_test.go index 1f06d8d0..61db3cba 100644 --- a/cmd/nvidia-ctk/cdi/generate/generate_test.go +++ b/cmd/nvidia-ctk/cdi/generate/generate_test.go @@ -37,7 +37,9 @@ func TestGenerateSpec(t *testing.T) { require.NoError(t, err) driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1") - disableHook := cli.NewStringSlice("enable-cuda-compat", "create-symlinks", "update-ldcache") + disableHook1 := cli.NewStringSlice("enable-cuda-compat") + disableHook2 := cli.NewStringSlice("enable-cuda-compat", "update-ldcache") + disableHook3 := cli.NewStringSlice("all") logger, _ := testlog.NewNullLogger() testCases := []struct { @@ -117,14 +119,14 @@ containerEdits: `, }, { - description: "disableHooks", + description: "disableHooks1", options: options{ - format: "yaml", - mode: "nvml", - vendor: "example.com", - class: "device", - driverRoot: driverRoot, - disableHooks: *disableHook, + format: "yaml", + mode: "nvml", + vendor: "example.com", + class: "device", + driverRoot: driverRoot, + disabledHooks: *disableHook1, }, expectedOptions: options{ format: "yaml", @@ -133,7 +135,128 @@ containerEdits: class: "device", nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook", driverRoot: driverRoot, - disableHooks: *disableHook, + disabledHooks: *disableHook1, + }, + expectedSpec: `--- +cdiVersion: 0.5.0 +kind: example.com/device +devices: + - name: "0" + containerEdits: + deviceNodes: + - path: /dev/nvidia0 + hostPath: {{ .driverRoot }}/dev/nvidia0 + - name: all + containerEdits: + deviceNodes: + - path: /dev/nvidia0 + hostPath: {{ .driverRoot }}/dev/nvidia0 +containerEdits: + env: + - NVIDIA_VISIBLE_DEVICES=void + deviceNodes: + - path: /dev/nvidiactl + hostPath: {{ .driverRoot }}/dev/nvidiactl + hooks: + - hookName: createContainer + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - create-symlinks + - --link + - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so + - hookName: createContainer + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - update-ldcache + - --folder + - /lib/x86_64-linux-gnu + mounts: + - hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77 + containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77 + options: + - ro + - nosuid + - nodev + - bind +`, + }, + { + description: "disableHooks2", + options: options{ + format: "yaml", + mode: "nvml", + vendor: "example.com", + class: "device", + driverRoot: driverRoot, + disabledHooks: *disableHook2, + }, + expectedOptions: options{ + format: "yaml", + mode: "nvml", + vendor: "example.com", + class: "device", + nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook", + driverRoot: driverRoot, + disabledHooks: *disableHook2, + }, + expectedSpec: `--- +cdiVersion: 0.5.0 +kind: example.com/device +devices: + - name: "0" + containerEdits: + deviceNodes: + - path: /dev/nvidia0 + hostPath: {{ .driverRoot }}/dev/nvidia0 + - name: all + containerEdits: + deviceNodes: + - path: /dev/nvidia0 + hostPath: {{ .driverRoot }}/dev/nvidia0 +containerEdits: + env: + - NVIDIA_VISIBLE_DEVICES=void + deviceNodes: + - path: /dev/nvidiactl + hostPath: {{ .driverRoot }}/dev/nvidiactl + hooks: + - hookName: createContainer + path: /usr/bin/nvidia-cdi-hook + args: + - nvidia-cdi-hook + - create-symlinks + - --link + - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so + mounts: + - hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77 + containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77 + options: + - ro + - nosuid + - nodev + - bind +`, + }, + { + description: "disableHooksAll", + options: options{ + format: "yaml", + mode: "nvml", + vendor: "example.com", + class: "device", + driverRoot: driverRoot, + disabledHooks: *disableHook3, + }, + expectedOptions: options{ + format: "yaml", + mode: "nvml", + vendor: "example.com", + class: "device", + nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook", + driverRoot: driverRoot, + disabledHooks: *disableHook3, }, expectedSpec: `--- cdiVersion: 0.5.0 diff --git a/pkg/nvcdi/api.go b/pkg/nvcdi/api.go index 2988026f..8b0fd9c8 100644 --- a/pkg/nvcdi/api.go +++ b/pkg/nvcdi/api.go @@ -44,4 +44,17 @@ const ( // HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility. // This was added with v1.17.5 of the NVIDIA Container Toolkit. HookEnableCudaCompat = HookName("enable-cuda-compat") + // HookCreateSymlinks refers to the hook used create symlinks inside the + // directory path to be mounted into a container. + HookCreateSymlinks = HookName("create-symlinks") + // HookUpdateLDCache refers to the hook used to Update the dynamic linker + // cache inside the directory path to be mounted into a container. + HookUpdateLDCache = HookName("update-ldcache") ) + +// AllHooks maintains a future-proof list of all defined hooks. +var AllHooks = []HookName{ + HookEnableCudaCompat, + HookCreateSymlinks, + HookUpdateLDCache, +} diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go index d3706e5c..145209e5 100644 --- a/pkg/nvcdi/driver-nvml.go +++ b/pkg/nvcdi/driver-nvml.go @@ -99,7 +99,7 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover var discoverers []discover.Discover - if l.HookIsSupported("create-symlinks") { + if l.HookIsSupported(HookCreateSymlinks) { driverDotSoSymlinksDiscoverer := discover.WithDriverDotSoSymlinks( libraries, version, @@ -118,7 +118,7 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover discoverers = append(discoverers, cudaCompatLibHookDiscoverer) } - if l.HookIsSupported("update-ldcache") { + if l.HookIsSupported(HookUpdateLDCache) { updateLDCache, _ := discover.NewLDCacheUpdateHook( l.logger, libraries, diff --git a/pkg/nvcdi/hooks.go b/pkg/nvcdi/hooks.go index a4620dc8..dfc453f0 100644 --- a/pkg/nvcdi/hooks.go +++ b/pkg/nvcdi/hooks.go @@ -28,3 +28,14 @@ func (l *nvcdilib) HookIsSupported(h HookName) bool { } return !l.disabledHooks[h] } + +func (d disabledHooks) Set(value HookName) { + if value == "all" { + for _, hook := range AllHooks { + d[hook] = true + } + return + } + + d[value] = true +} diff --git a/pkg/nvcdi/options.go b/pkg/nvcdi/options.go index f38f2b4a..514391e3 100644 --- a/pkg/nvcdi/options.go +++ b/pkg/nvcdi/options.go @@ -158,11 +158,11 @@ func WithLibrarySearchPaths(paths []string) Option { // WithDisabledHook allows specific hooks to the disabled. // This option can be specified multiple times for each hook. -func WithDisabledHook(hook HookName) Option { +func WithDisabledHook[T string | HookName](hook T) Option { return func(o *nvcdilib) { if o.disabledHooks == nil { o.disabledHooks = make(map[HookName]bool) } - o.disabledHooks[hook] = true + o.disabledHooks.Set(HookName(hook)) } }