From 96ecfd0b6f7458de978151c6c9f8476b47d9c9d9 Mon Sep 17 00:00:00 2001
From: Carlos Eduardo Arango Gutierrez <eduardoa@nvidia.com>
Date: Mon, 12 May 2025 18:52:37 +0200
Subject: [PATCH] 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 <eduardoa@nvidia.com>
---
 cmd/nvidia-ctk/cdi/generate/generate.go      | 19 +++++-
 cmd/nvidia-ctk/cdi/generate/generate_test.go | 66 ++++++++++++++++++++
 2 files changed, 83 insertions(+), 2 deletions(-)

diff --git a/cmd/nvidia-ctk/cdi/generate/generate.go b/cmd/nvidia-ctk/cdi/generate/generate.go
index b187335b..14fa6af0 100644
--- 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
+	disableHooks       cli.StringSlice
 
 	csv struct {
 		files          cli.StringSlice
@@ -176,6 +177,12 @@ func (m command) build() *cli.Command {
 			Usage:       "Specify a pattern the CSV mount specifications.",
 			Destination: &opts.csv.ignorePatterns,
 		},
+		&cli.StringSliceFlag{
+			Name:        "disable-hook",
+			Usage:       "Comma-separated list of hooks to skip when generating the CDI specification.",
+			Value:       cli.NewStringSlice(),
+			Destination: &opts.disableHooks,
+		},
 	}
 
 	return &c
@@ -262,7 +269,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 +283,15 @@ 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),
-	)
+	}
+
+	if len(opts.disableHooks.Value()) > 0 {
+		for _, hook := range opts.disableHooks.Value() {
+			initOpts = append(initOpts, nvcdi.WithDisabledHook(nvcdi.HookName(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
index 9e9a6a4e..ed011c13 100644
--- 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,7 @@ func TestGenerateSpec(t *testing.T) {
 	require.NoError(t, err)
 
 	driverRoot := filepath.Join(moduleRoot, "testdata", "lookup", "rootfs-1")
+	disableHook := cli.NewStringSlice("enable-cuda-compat")
 
 	logger, _ := testlog.NewNullLogger()
 	testCases := []struct {
@@ -112,6 +114,70 @@ containerEdits:
             - nosuid
             - nodev
             - bind
+`,
+		},
+		{
+			description: "skipHook",
+			options: options{
+				format:       "yaml",
+				mode:         "nvml",
+				vendor:       "example.com",
+				class:        "device",
+				driverRoot:   driverRoot,
+				disableHooks: *disableHook,
+			},
+			expectedOptions: options{
+				format:            "yaml",
+				mode:              "nvml",
+				vendor:            "example.com",
+				class:             "device",
+				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
+				driverRoot:        driverRoot,
+				disableHooks:      *disableHook,
+			},
+			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
 `,
 		},
 	}