Use require-nvidia-kernel-modules feature for toolkit installation

Signed-off-by: Evan Lezar <[email protected]>
NVIDIA · Oct 30, 2024 · bf34b13 · bf34b13
1 parent 7263d26
commit bf34b13
Show file tree

Hide file tree

Showing 5 changed files with 5 additions and 44 deletions.
diff --git a/tools/container/toolkit/executable.go b/tools/container/toolkit/executable.go
@@ -36,7 +36,6 @@ type executable struct {
 	source   string
 	target   executableTarget
 	env      map[string]string
-	preLines []string
 	argLines []string
 }
 
@@ -96,11 +95,6 @@ func (e executable) writeWrapperTo(wrapper io.Writer, destFolder string, dotfile
 	// Add the shebang
 	fmt.Fprintln(wrapper, "#! /bin/sh")
 
-	// Add the preceding lines if any
-	for _, line := range e.preLines {
-		fmt.Fprintf(wrapper, "%s\n", r.apply(line))
-	}
-
 	// Update the path to include the destination folder
 	var env map[string]string
 	if e.env == nil {

diff --git a/tools/container/toolkit/executable_test.go b/tools/container/toolkit/executable_test.go
@@ -59,23 +59,6 @@ func TestWrapper(t *testing.T) {
 				"",
 			},
 		},
-		{
-			e: executable{
-				preLines: []string{
-					"preline1",
-					"preline2",
-				},
-			},
-			expectedLines: []string{
-				shebang,
-				"preline1",
-				"preline2",
-				"PATH=/dest/folder:$PATH \\",
-				"source.real \\",
-				"\t\"$@\"",
-				"",
-			},
-		},
 		{
 			e: executable{
 				argLines: []string{

diff --git a/tools/container/toolkit/runtime.go b/tools/container/toolkit/runtime.go
@@ -57,27 +57,16 @@ func newNvidiaContainerRuntimeInstaller(source string) *executable {
 }
 
 func newRuntimeInstaller(source string, target executableTarget, env map[string]string) *executable {
-	preLines := []string{
-		"",
-		"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
-		"if [ \"${?}\" != \"0\" ]; then",
-		"	echo \"nvidia driver modules are not yet loaded, invoking runc directly\"",
-		"	exec runc \"$@\"",
-		"fi",
-		"",
-	}
-
 	runtimeEnv := make(map[string]string)
 	runtimeEnv["XDG_CONFIG_HOME"] = filepath.Join(destDirPattern, ".config")
 	for k, v := range env {
 		runtimeEnv[k] = v
 	}
 
 	r := executable{
-		source:   source,
-		target:   target,
-		env:      runtimeEnv,
-		preLines: preLines,
+		source: source,
+		target: target,
+		env:    runtimeEnv,
 	}
 
 	return &r

diff --git a/tools/container/toolkit/runtime_test.go b/tools/container/toolkit/runtime_test.go
@@ -38,13 +38,6 @@ func TestNvidiaContainerRuntimeInstallerWrapper(t *testing.T) {
 
 	expectedLines := []string{
 		shebang,
-		"",
-		"cat /proc/modules | grep -e \"^nvidia \" >/dev/null 2>&1",
-		"if [ \"${?}\" != \"0\" ]; then",
-		"	echo \"nvidia driver modules are not yet loaded, invoking runc directly\"",
-		"	exec runc \"$@\"",
-		"fi",
-		"",
 		"PATH=/dest/folder:$PATH \\",
 		"XDG_CONFIG_HOME=/dest/folder/.config \\",
 		"source.real \\",

diff --git a/tools/container/toolkit/toolkit.go b/tools/container/toolkit/toolkit.go
@@ -458,6 +458,8 @@ func installToolkitConfig(c *cli.Context, toolkitConfigPath string, nvidiaContai
 		// Set the nvidia-container-runtime-hook options
 		"nvidia-container-runtime-hook.path":                nvidaContainerRuntimeHookPath,
 		"nvidia-container-runtime-hook.skip-mode-detection": opts.ContainerRuntimeHookSkipModeDetection,
+		// We require the NVIDIA kernel modules to be loaded.
+		"features.require-nvidia-kernel-modules": true,
 	}
 
 	toolkitRuntimeList := opts.ContainerRuntimeRuntimes.Value()