From 047dd5bef07d0ad9ea89b12c845e7c7e96b27c1d Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Thu, 6 Feb 2025 23:17:21 +0100
Subject: [PATCH] Add cuda-compat hook to CDI spec generation

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go | 6 ++++++
 cmd/nvidia-ctk/cdi/generate/generate_test.go               | 6 ++++++
 pkg/nvcdi/driver-nvml.go                                   | 3 +++
 3 files changed, 15 insertions(+)

diff --git a/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go b/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go
index 855141ff..5a28a463 100644
--- a/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go
+++ b/cmd/nvidia-ctk-installer/container/toolkit/toolkit_test.go
@@ -80,6 +80,12 @@ containerEdits:
     - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
     hookName: createContainer
     path: {{ .toolkitRoot }}/nvidia-cdi-hook
+  - args:
+    - nvidia-cdi-hook
+    - cuda-compat
+    - --host-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
index b428dca3..12e73d65 100644
--- a/cmd/nvidia-ctk/cdi/generate/generate_test.go
+++ b/cmd/nvidia-ctk/cdi/generate/generate_test.go
@@ -79,6 +79,12 @@ containerEdits:
     - libcuda.so.1::/lib/x86_64-linux-gnu/libcuda.so
     hookName: createContainer
     path: /usr/bin/nvidia-cdi-hook
+  - args:
+    - nvidia-cdi-hook
+    - cuda-compat
+    - --host-driver-version=999.88.77
+    hookName: createContainer
+    path: /usr/bin/nvidia-cdi-hook
   - args:
     - nvidia-cdi-hook
     - update-ldcache
diff --git a/pkg/nvcdi/driver-nvml.go b/pkg/nvcdi/driver-nvml.go
index b0006aeb..782d60fc 100644
--- 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.
+	cudaCompatLibHookDiscoverer := discover.NewCUDACompatHookDiscoverer(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,
 		),
+		cudaCompatLibHookDiscoverer,
 		updateLDCache,
 	)