Add compat-lib hook to CDI spec generation

Signed-off-by: Evan Lezar <[email protected]>
NVIDIA · Feb 6, 2025 · ba30397 · ba30397
1 parent ed24e15
commit ba30397
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 5 deletions.
diff --git a/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go b/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go
@@ -80,6 +80,13 @@ containerEdits:
     - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
     hookName: createContainer
     path: {{ .toolkitRoot }}/nvidia-cdi-hook
+  - args:
+    - nvidia-cdi-hook
+    - compat-libs
+    - --driver-version
+    - 999.88.77
+    hookName: createContainer
+    path: {{ .toolkitRoot }}/nvidia-cdi-hook
   - args:
     - nvidia-cdi-hook
     - update-ldcache

diff --git a/cmd/nvidia-ctk/cdi/generate/generate_test.go b/cmd/nvidia-ctk/cdi/generate/generate_test.go
@@ -19,6 +19,7 @@ package generate
 import (
 	"bytes"
 	"path/filepath"
+	"strings"
 	"testing"
 
 	"github.com/NVIDIA/go-nvml/pkg/nvml"
@@ -66,7 +67,7 @@ func TestGenerateSpec(t *testing.T) {
 cdiVersion: 0.5.0
 containerEdits:
   deviceNodes:
-  - hostPath: /Users/elezar/dev/container-toolkit/testdata/lookup/rootfs-1/dev/nvidiactl
+  - hostPath: {{ .driverRoot }}/dev/nvidiactl
     path: /dev/nvidiactl
   env:
   - NVIDIA_VISIBLE_DEVICES=void
@@ -78,6 +79,13 @@ containerEdits:
     - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
     hookName: createContainer
     path: /usr/bin/nvidia-cdi-hook
+  - args:
+    - nvidia-cdi-hook
+    - compat-libs
+    - --driver-version
+    - 999.88.77
+    hookName: createContainer
+    path: /usr/bin/nvidia-cdi-hook
   - args:
     - nvidia-cdi-hook
     - update-ldcache
@@ -87,7 +95,7 @@ containerEdits:
     path: /usr/bin/nvidia-cdi-hook
   mounts:
   - containerPath: /lib/x86_64-linux-gnu/libcuda.so.999.88.77
-    hostPath: /Users/elezar/dev/container-toolkit/testdata/lookup/rootfs-1/lib/x86_64-linux-gnu/libcuda.so.999.88.77
+    hostPath: {{ .driverRoot }}/lib/x86_64-linux-gnu/libcuda.so.999.88.77
     options:
     - ro
     - nosuid
@@ -96,12 +104,12 @@ containerEdits:
 devices:
 - containerEdits:
     deviceNodes:
-    - hostPath: /Users/elezar/dev/container-toolkit/testdata/lookup/rootfs-1/dev/nvidia0
+    - hostPath: {{ .driverRoot }}/dev/nvidia0
       path: /dev/nvidia0
   name: "0"
 - containerEdits:
     deviceNodes:
-    - hostPath: /Users/elezar/dev/container-toolkit/testdata/lookup/rootfs-1/dev/nvidia0
+    - hostPath: {{ .driverRoot }}/dev/nvidia0
       path: /dev/nvidia0
   name: all
 kind: example.com/device
@@ -144,7 +152,7 @@ kind: example.com/device
 			_, err = spec.WriteTo(&buf)
 			require.NoError(t, err)
 
-			require.Equal(t, tc.expectedSpec, buf.String())
+			require.Equal(t, strings.ReplaceAll(tc.expectedSpec, "{{ .driverRoot }}", driverRoot), buf.String())
 		})
 	}
 }
diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go
@@ -97,6 +97,8 @@ func NewDriverLibraryDiscoverer(logger logger.Interface, driver *root.Driver, nv
 		libraryPaths,
 	)
 
+	// TODO: The following should use the version directly.
+	compatLibHookDiscoverer := discover.NewCUDACompatLibsDiscoverer(logger, nvidiaCDIHookPath, driver)
 	updateLDCache, _ := discover.NewLDCacheUpdateHook(logger, libraries, nvidiaCDIHookPath, ldconfigPath)
 
 	d := discover.Merge(
@@ -105,6 +107,7 @@ func NewDriverLibraryDiscoverer(logger logger.Interface, driver *root.Driver, nv
 			version,
 			nvidiaCDIHookPath,
 		),
+		compatLibHookDiscoverer,
 		updateLDCache,
 	)