From bf34b137c1acd3051b1dd91fe671c84d7a284d87 Mon Sep 17 00:00:00 2001
From: Evan Lezar <elezar@nvidia.com>
Date: Wed, 30 Oct 2024 15:13:34 +0100
Subject: [PATCH] Use require-nvidia-kernel-modules feature for toolkit
 installation

Signed-off-by: Evan Lezar <elezar@nvidia.com>
---
 tools/container/toolkit/executable.go      |  6 ------
 tools/container/toolkit/executable_test.go | 17 -----------------
 tools/container/toolkit/runtime.go         | 17 +++--------------
 tools/container/toolkit/runtime_test.go    |  7 -------
 tools/container/toolkit/toolkit.go         |  2 ++
 5 files changed, 5 insertions(+), 44 deletions(-)

diff --git a/tools/container/toolkit/executable.go b/tools/container/toolkit/executable.go
index 394ca0076..7fb879b70 100644
--- 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
index 8cb47596c..5b3249ea7 100644
--- 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
index bdfca9834..1f3ce31bc 100644
--- a/tools/container/toolkit/runtime.go
+++ b/tools/container/toolkit/runtime.go
@@ -57,16 +57,6 @@ 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 {
@@ -74,10 +64,9 @@ func newRuntimeInstaller(source string, target executableTarget, env map[string]
 	}
 
 	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
index d2841506d..77d84ec7c 100644
--- 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
index 43e68ca55..2580740e9 100644
--- 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()