From 2d9e4cb720597e666bc82294515bfe737d4e5dae Mon Sep 17 00:00:00 2001
From: Evan Lezar <7723350-elezar@users.noreply.gitlab.com>
Date: Thu, 9 Mar 2023 16:44:13 +0000
Subject: [PATCH] Merge branch 'set-nvidia-ctk-path' into 'main'

Set nvidia-ctk.path config option based on installed path

See merge request nvidia/container-toolkit/container-toolkit!334
---
 test/container/toolkit_test.sh     |  1 +
 tools/container/toolkit/toolkit.go | 17 ++++++++++-------
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/test/container/toolkit_test.sh b/test/container/toolkit_test.sh
index 84df4cac..617d4798 100644
--- a/test/container/toolkit_test.sh
+++ b/test/container/toolkit_test.sh
@@ -61,6 +61,7 @@ testing::toolkit::install() {
 	grep -q -E "^\s*ldconfig = \"@${nvidia_run_dir}/driver/sbin/ldconfig(.real)?\"" "${shared_dir}/usr/local/nvidia/toolkit/.config/nvidia-container-runtime/config.toml"
 	grep -q -E "^\s*root = \"${nvidia_run_dir}/driver\"" "${shared_dir}/usr/local/nvidia/toolkit/.config/nvidia-container-runtime/config.toml"
 	grep -q -E "^\s*path = \"/usr/local/nvidia/toolkit/nvidia-container-cli\"" "${shared_dir}/usr/local/nvidia/toolkit/.config/nvidia-container-runtime/config.toml"
+	grep -q -E "^\s*path = \"/usr/local/nvidia/toolkit/nvidia-ctk\"" "${shared_dir}/usr/local/nvidia/toolkit/.config/nvidia-container-runtime/config.toml"
 }
 
 testing::toolkit::delete() {
diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go
index abaeded6..a4e89b38 100644
--- a/tools/container/toolkit/toolkit.go
+++ b/tools/container/toolkit/toolkit.go
@@ -204,16 +204,16 @@ func Install(cli *cli.Context, opts *options) error {
 		return fmt.Errorf("error installing NVIDIA container runtime hook: %v", err)
 	}
 
-	err = installToolkitConfig(toolkitConfigPath, nvidiaContainerCliExecutable, opts)
-	if err != nil {
-		return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err)
-	}
-
-	_, err = installContainerToolkitCLI(opts.toolkitRoot)
+	nvidiaCTKPath, err := installContainerToolkitCLI(opts.toolkitRoot)
 	if err != nil {
 		return fmt.Errorf("error installing NVIDIA Container Toolkit CLI: %v", err)
 	}
 
+	err = installToolkitConfig(toolkitConfigPath, nvidiaContainerCliExecutable, nvidiaCTKPath, opts)
+	if err != nil {
+		return fmt.Errorf("error installing NVIDIA container toolkit config: %v", err)
+	}
+
 	return nil
 }
 
@@ -267,7 +267,7 @@ func installLibrary(libName string, toolkitRoot string) error {
 
 // installToolkitConfig installs the config file for the NVIDIA container toolkit ensuring
 // that the settings are updated to match the desired install and nvidia driver directories.
-func installToolkitConfig(toolkitConfigPath string, nvidiaContainerCliExecutablePath string, opts *options) error {
+func installToolkitConfig(toolkitConfigPath string, nvidiaContainerCliExecutablePath string, nvidiaCTKPath string, opts *options) error {
 	log.Infof("Installing NVIDIA container toolkit config '%v'", toolkitConfigPath)
 
 	config, err := toml.LoadFile(nvidiaContainerToolkitConfigSource)
@@ -316,6 +316,9 @@ func installToolkitConfig(toolkitConfigPath string, nvidiaContainerCliExecutable
 		config.Set(key, value)
 	}
 
+	// Set nvidia-ctk options
+	config.Set("nvidia-ctk.path", nvidiaCTKPath)
+
 	_, err = config.WriteTo(targetConfig)
 	if err != nil {
 		return fmt.Errorf("error writing config: %v", err)