Add --disable-hook flag to cdi generate command

ArangoGutierrez · ArangoGutierrez · commit 9aa88eaa7691 · 2025-05-23T11:08:09.000+02:00
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 &lt;eduardoa@nvidia.com&gt;
diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
@@ -57,6 +57,7 @@ type options struct {
 
 	configSearchPaths  cli.StringSlice
 	librarySearchPaths cli.StringSlice
+	disabledHooks      cli.StringSlice
 
 	csv struct {
 		files          cli.StringSlice
@@ -176,6 +177,13 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify a pattern the CSV mount specifications.",
 			Destination: &opts.csv.ignorePatterns,
 		},
+		&cli.StringSliceFlag{
+			Name:        "disable-hook",
+			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.disabledHooks,
+		},
 	}
 
 	return &c
@@ -262,7 +270,7 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		deviceNamers = append(deviceNamers, deviceNamer)
 	}
 
-	cdilib, err := nvcdi.New(
+	initOpts := []nvcdi.Option{
 		nvcdi.WithLogger(m.logger),
 		nvcdi.WithDriverRoot(opts.driverRoot),
 		nvcdi.WithDevRoot(opts.devRoot),
@@ -276,7 +284,13 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 		nvcdi.WithCSVIgnorePatterns(opts.csv.ignorePatterns.Value()),
 		// We set the following to allow for dependency injection:
 		nvcdi.WithNvmlLib(opts.nvmllib),
-	)
+	}
+
+	for _, hook := range opts.disabledHooks.Value() {
+		initOpts = append(initOpts, nvcdi.WithDisabledHook(hook))
+	}
+
+	cdilib, err := nvcdi.New(initOpts...)
 	if err != nil {
 		return nil, fmt.Errorf("failed to create CDI library: %v", err)
 	}
diff --git a/cmd/nvidia-ctk/cdi/generate/generate_test.go b/cmd/nvidia-ctk/cdi/generate/generate_test.go
@@ -26,6 +26,7 @@ import (
 	"github.com/NVIDIA/go-nvml/pkg/nvml/mock/dgxa100"
 	testlog "github.com/sirupsen/logrus/hooks/test"
 	"github.com/stretchr/testify/require"
+	"github.com/urfave/cli/v2"
 
 	"github.com/NVIDIA/nvidia-container-toolkit/internal/test"
 )
@@ -36,6 +37,9 @@ func TestGenerateSpec(t *testing.T) {
 	require.NoError(t, err)
 
 	driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")
+	disableHook1 := cli.NewStringSlice("enable-cuda-compat")
+	disableHook2 := cli.NewStringSlice("enable-cuda-compat", "update-ldcache")
+	disableHook3 := cli.NewStringSlice("all")
 
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
@@ -113,6 +117,179 @@ containerEdits:
             - nodev
             - rbind
             - rprivate
+`,
+		},
+		{
+			description: "disableHooks1",
+			options: options{
+				format:        "yaml",
+				mode:          "nvml",
+				vendor:        "example.com",
+				class:         "device",
+				driverRoot:    driverRoot,
+				disabledHooks: *disableHook1,
+			},
+			expectedOptions: options{
+				format:            "yaml",
+				mode:              "nvml",
+				vendor:            "example.com",
+				class:             "device",
+				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
+				driverRoot:        driverRoot,
+				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
+            - rbind
+            - rprivate
+`,
+		},
+		{
+			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
+            - rbind
+            - rprivate
+`,
+		},
+		{
+			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
+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
+    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
+            - rbind
+            - rprivate
 `,
 		},
 	}
diff --git a/internal/discover/hooks.go b/internal/discover/hooks.go
@@ -46,8 +46,16 @@ func (h *Hook) Hooks() ([]Hook, error) {
 
 type HookName string
 
-// DisabledHooks allows individual hooks to be disabled.
-type DisabledHooks map[HookName]bool
+// disabledHooks allows individual hooks to be disabled.
+type disabledHooks map[HookName]bool
+
+// isDisabled checks if a hook is disabled.
+func (d disabledHooks) isDisabled(h HookName) bool {
+	if d["all"] {
+		return true
+	}
+	return d[h]
+}
 
 const (
 	// HookEnableCudaCompat refers to the hook used to enable CUDA Forward Compatibility.
@@ -67,26 +75,43 @@ var AllHooks = []HookName{
 	HookUpdateLDCache,
 }
 
-// Option is a function that configures the nvcdilib
 type Option func(*CDIHook)
 
 type CDIHook struct {
 	nvidiaCDIHookPath string
+	disabledHooks     disabledHooks
 }
 
 type HookCreator interface {
 	Create(HookName, ...string) *Hook
 }
 
-func NewHookCreator(nvidiaCDIHookPath string) HookCreator {
+func WithDisabledHooks(hooks ...HookName) Option {
+	return func(c *CDIHook) {
+		for _, hook := range hooks {
+			c.disabledHooks[hook] = true
+		}
+	}
+}
+
+func NewHookCreator(nvidiaCDIHookPath string, opts ...Option) HookCreator {
 	CDIHook := &CDIHook{
 		nvidiaCDIHookPath: nvidiaCDIHookPath,
+		disabledHooks:     disabledHooks{},
+	}
+
+	for _, opt := range opts {
+		opt(CDIHook)
 	}
 
 	return CDIHook
 }
 
 func (c CDIHook) Create(name HookName, args ...string) *Hook {
+	if c.disabledHooks[name] {
+		return nil
+	}
+
 	if name == "create-symlinks" {
 		if len(args) == 0 {
 			return nil
diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go
@@ -106,11 +106,9 @@ func (l *nvcdilib) NewDriverLibraryDiscoverer(version string) (discover.Discover
 	)
 	discoverers = append(discoverers, driverDotSoSymlinksDiscoverer)
 
-	if l.HookIsSupported(HookEnableCudaCompat) {
-		// TODO: The following should use the version directly.
-		cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, l.driver)
-		discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
-	}
+	// TODO: The following should use the version directly.
+	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(l.logger, l.hookCreator, l.driver)
+	discoverers = append(discoverers, cudaCompatLibHookDiscoverer)
 
 	updateLDCache, _ := discover.NewLDCacheUpdateHook(l.logger, libraries, l.hookCreator, l.ldconfigPath)
 	discoverers = append(discoverers, updateLDCache)
diff --git a/pkg/nvcdi/hooks.go b/pkg/nvcdi/hooks.go
diff --git a/pkg/nvcdi/lib.go b/pkg/nvcdi/lib.go
@@ -56,15 +56,13 @@ type nvcdilib struct {
 
 	mergedDeviceOptions []transform.MergedDeviceOption
 
-	disabledHooks discover.DisabledHooks
+	disabledHooks []discover.HookName
 	hookCreator   discover.HookCreator
 }
 
 // New creates a new nvcdi library
 func New(opts ...Option) (Interface, error) {
-	l := &nvcdilib{
-		disabledHooks: make(discover.DisabledHooks),
-	}
+	l := &nvcdilib{}
 	for _, opt := range opts {
 		opt(l)
 	}
@@ -81,8 +79,6 @@ func New(opts ...Option) (Interface, error) {
 	if l.nvidiaCDIHookPath == "" {
 		l.nvidiaCDIHookPath = "/usr/bin/nvidia-cdi-hook"
 	}
-	// create hookCreator
-	l.hookCreator = discover.NewHookCreator(l.nvidiaCDIHookPath)
 
 	if l.driverRoot == "" {
 		l.driverRoot = "/"
@@ -150,7 +146,7 @@ func New(opts ...Option) (Interface, error) {
 			l.vendor = "management.nvidia.com"
 		}
 		// Management containers in general do not require CUDA Forward compatibility.
-		l.disabledHooks[HookEnableCudaCompat] = true
+		l.disabledHooks = append(l.disabledHooks, HookEnableCudaCompat)
 		lib = (*managementlib)(l)
 	case ModeNvml:
 		lib = (*nvmllib)(l)
@@ -175,6 +171,9 @@ func New(opts ...Option) (Interface, error) {
 		return nil, fmt.Errorf("unknown mode %q", l.mode)
 	}
 
+	// create hookCreator
+	l.hookCreator = discover.NewHookCreator(l.nvidiaCDIHookPath, discover.WithDisabledHooks(l.disabledHooks...))
+
 	w := wrapper{
 		Interface:           lib,
 		vendor:              l.vendor,
diff --git a/pkg/nvcdi/options.go b/pkg/nvcdi/options.go
