diff --git a/test/container/toolkit_test.sh b/test/container/toolkit_test.sh index f92bb0ff..6ae724ab 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 80f81737..5bae0f5c 100644 --- a/tools/container/toolkit/toolkit.go +++ b/tools/container/toolkit/toolkit.go @@ -257,16 +257,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) - } - 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 generateCDISpec(opts, nvidiaCTKPath) } @@ -320,7 +320,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) @@ -368,6 +368,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)