From b9b9b5f472b4ad60a7f9a6a7d58ad9a9632efdac Mon Sep 17 00:00:00 2001
From: Vitaliy Emporopulo <vemporop@redhat.com>
Date: Mon, 12 Feb 2024 15:13:10 +0200
Subject: [PATCH] Allow custom NVIDIA CTK path

Signed-off-by: Vitaliy Emporopulo <vemporop@redhat.com>
---
 cmd/nvidia-dra-plugin/cdi.go                                 | 1 +
 deployments/helm/k8s-dra-driver/templates/kubeletplugin.yaml | 2 ++
 deployments/helm/k8s-dra-driver/values.yaml                  | 5 +++++
 3 files changed, 8 insertions(+)

diff --git a/cmd/nvidia-dra-plugin/cdi.go b/cmd/nvidia-dra-plugin/cdi.go
index 5c405e549..997865fbe 100644
--- a/cmd/nvidia-dra-plugin/cdi.go
+++ b/cmd/nvidia-dra-plugin/cdi.go
@@ -91,6 +91,7 @@ func NewCDIHandler(opts ...cdiOption) (*CDIHandler, error) {
 			nvcdi.WithMode("nvml"),
 			nvcdi.WithVendor(h.vendor),
 			nvcdi.WithClass(h.class),
+			nvcdi.WithNVIDIACTKPath(h.nvidiaCTKPath),
 		)
 		if err != nil {
 			return nil, fmt.Errorf("unable to create CDI library: %w", err)
diff --git a/deployments/helm/k8s-dra-driver/templates/kubeletplugin.yaml b/deployments/helm/k8s-dra-driver/templates/kubeletplugin.yaml
index c2b7dd69c..bb0a4a572 100644
--- a/deployments/helm/k8s-dra-driver/templates/kubeletplugin.yaml
+++ b/deployments/helm/k8s-dra-driver/templates/kubeletplugin.yaml
@@ -76,6 +76,8 @@ spec:
         resources:
           {{- toYaml .Values.kubeletPlugin.containers.plugin.resources | nindent 10 }}
         env:
+        - name: NVIDIA_CTK_PATH
+          value: "{{ .Values.nvidiaCtkPath }}"
         - name: NVIDIA_DRIVER_ROOT
           value: "{{ .Values.nvidiaDriverRoot }}"
         - name: NVIDIA_VISIBLE_DEVICES
diff --git a/deployments/helm/k8s-dra-driver/values.yaml b/deployments/helm/k8s-dra-driver/values.yaml
index 41d77cf12..512cdd7c9 100644
--- a/deployments/helm/k8s-dra-driver/values.yaml
+++ b/deployments/helm/k8s-dra-driver/values.yaml
@@ -22,6 +22,11 @@
 # For driver installed directly on a host, a value of `/` is used.
 nvidiaDriverRoot: /
 
+# Specify the path of CTK binary (nvidia-ctk) on the host,
+# as it should appear in the the generated CDI specification.
+# The path depends on the system that runs on the node.
+nvidiaCtkPath: /usr/bin/nvidia-ctk
+
 nameOverride: ""
 fullnameOverride: ""
 namespaceOverride: ""