fixup! USE Matchers

Signed-off-by: Evan Lezar <[email protected]>

Author: elezar

tests/e2e/nvidia-cdi-refresh_test.go

@@ -134,7 +134,7 @@ var _ = Describe("nvidia-cdi-refresh", Ordered, ContinueOnFailure, Label("system
 		var err error
 		// TODO: We set installCTK to true here to SKIP the mounting of the files from the host.
 		// The test here does NOT require the host toolkit.
-		systemdRunner, err = NewNestedContainerRunner(runner, outerContainerImage, false, containerName, localCacheDir)
+		systemdRunner, err = NewNestedContainerRunner(runner, outerContainerImage, false, containerName, localCacheDir, true)
 		Expect(err).ToNot(HaveOccurred())
 		for range 10 {
 			state, _, err := systemdRunner.Run(getSystemStateScript)

tests/e2e/nvidia-container-cli_test.go

@@ -78,7 +78,7 @@ var _ = Describe("nvidia-container-cli", Ordered, ContinueOnFailure, Label("libn
 
 	BeforeAll(func(ctx context.Context) {
 		var err error
-		nestedContainerRunner, err = NewNestedContainerRunner(runner, "ubuntu", !installCTK, containerName, localCacheDir)
+		nestedContainerRunner, err = NewNestedContainerRunner(runner, "ubuntu", !installCTK, containerName, localCacheDir, true)
 		Expect(err).ToNot(HaveOccurred())
 
 		if installCTK {

tests/e2e/nvidia-ctk_containerd_test.go

@@ -121,7 +121,7 @@ var _ = Describe("containerd", Ordered, ContinueOnFailure, Label("container-runt
 
 				// Create the nested container with the global cache mounted
 				// TODO: This runner doesn't actually NEED GPU access.
-				nestedContainerRunner, err = NewNestedContainerRunner(runner, env.image, false, containerName, localCacheDir)
+				nestedContainerRunner, err = NewNestedContainerRunner(runner, env.image, false, containerName, localCacheDir, false)
 				Expect(err).ToNot(HaveOccurred())
 
 				// Backup original containerd configuration

tests/e2e/runner.go

@@ -100,7 +100,7 @@ func NewRunner(opts ...runnerOption) Runner {
 // NewNestedContainerRunner creates a new nested container runner.
 // A nested container runs a container inside another container based on a
 // given runner (remote or local).
-func NewNestedContainerRunner(runner Runner, baseImage string, mountToolkitFromHost bool, containerName string, cacheDir string) (Runner, error) {
+func NewNestedContainerRunner(runner Runner, baseImage string, mountToolkitFromHost bool, containerName string, cacheDir string, requiresGPUs bool) (Runner, error) {
 	// If a container with the same name exists from a previous test run, remove it first.
 	// Ignore errors as container might not exist
 	_, _, err := runner.Run(fmt.Sprintf("docker rm -f %s 2>/dev/null || true", containerName))
@@ -110,6 +110,16 @@ func NewNestedContainerRunner(runner Runner, baseImage string, mountToolkitFromH
 
 	var additionalContainerArguments []string
 
+	if requiresGPUs {
+		// If the container requires access to GPUs we explicitly add the nvidia
+		// runtime and set `NVIDIA_VISIBLE_DEVICES` to trigger jit-cdi spec
+		// generation.
+		additionalContainerArguments = append(additionalContainerArguments,
+			"--runtime=nvidia",
+			"-e NVIDIA_VISIBLE_DEVICES=runtime.nvidia.com/gpu=all",
+		)
+	}
+
 	if cacheDir != "" {
 		additionalContainerArguments = append(additionalContainerArguments,
 			"-v "+cacheDir+":"+cacheDir+":ro",
@@ -302,10 +312,6 @@ func connectOrDie(sshKey, sshUser, host, port string) (*ssh.Client, error) {
 
 // outerContainerTemplate represents a template to start a container with
 // a name specified.
-// The container is given access to all NVIDIA gpus by explicitly using the
-// nvidia runtime and the `runtime.nvidia.com/gpu=all` device to trigger JIT
-// CDI spec generation.
-// The template also allows for additional arguments to be specified.
 type outerContainer struct {
 	Name                string
 	BaseImage           string
@@ -314,8 +320,6 @@ type outerContainer struct {
 
 func (o *outerContainer) Render() (string, error) {
 	tmpl, err := template.New("startContainer").Parse(`docker run -d --name {{.Name}} --privileged \
--e NVIDIA_VISIBLE_DEVICES=runtime.nvidia.com/gpu=all \
--e NVIDIA_DRIVER_CAPABILITIES=all \
 {{ range $i, $a := .AdditionalArguments -}}
 {{ $a }} \
 {{ end -}}