Add all option to disable all hooks

Signed-off-by: Carlos Eduardo Arango Gutierrez <[email protected]>

Author: ArangoGutierrez

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

@@ -287,7 +287,9 @@ func (m command) generateSpec(opts *options) (spec.Interface, error) {
 
 	if len(opts.disableHooks.Value()) > 0 {
 		for _, hook := range opts.disableHooks.Value() {
-			initOpts = append(initOpts, nvcdi.WithDisabledHook(nvcdi.HookName(hook)))
+			for _, hookName := range nvcdi.NewHookName(hook) {
+				initOpts = append(initOpts, nvcdi.WithDisabledHook(hookName))
+			}
 		}
 	}
 

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,
+				disableHooks: *disableHook1,
 			},
 			expectedOptions: options{
 				format:            "yaml",
@@ -133,7 +135,128 @@ containerEdits:
 				class:             "device",
 				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
 				driverRoot:        driverRoot,
-				disableHooks:      *disableHook,
+				disableHooks:      *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,
+				disableHooks: *disableHook2,
+			},
+			expectedOptions: options{
+				format:            "yaml",
+				mode:              "nvml",
+				vendor:            "example.com",
+				class:             "device",
+				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
+				driverRoot:        driverRoot,
+				disableHooks:      *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,
+				disableHooks: *disableHook3,
+			},
+			expectedOptions: options{
+				format:            "yaml",
+				mode:              "nvml",
+				vendor:            "example.com",
+				class:             "device",
+				nvidiaCDIHookPath: "/usr/bin/nvidia-cdi-hook",
+				driverRoot:        driverRoot,
+				disableHooks:      *disableHook3,
 			},
 			expectedSpec: `---
 cdiVersion: 0.5.0

pkg/nvcdi/api.go

@@ -44,4 +44,33 @@ 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")
 )
+
+// NewHookName takes a string and returns a []HookName, empty if the HookName
+// is invalid and all Hooks if the string is "all"
+func NewHookName(hookName string) []HookName {
+	if hookName == "" {
+		return []HookName{}
+	}
+
+	if hookName == "all" {
+		return []HookName{HookEnableCudaCompat, HookCreateSymlinks, HookUpdateLDCache}
+	}
+
+	switch hookName {
+	case string(HookEnableCudaCompat):
+		return []HookName{HookEnableCudaCompat}
+	case string(HookCreateSymlinks):
+		return []HookName{HookCreateSymlinks}
+	case string(HookUpdateLDCache):
+		return []HookName{HookUpdateLDCache}
+	default:
+		return nil
+	}
+}

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,