Move GPU test utilities to its own package.

These will be reusable across GPU tests, not just the smoke tests.

PiperOrigin-RevId: 587888928

pkg/test/dockerutil/BUILD

@@ -12,6 +12,7 @@ go_library(
         "container.go",
         "dockerutil.go",
         "exec.go",
+        "gpu.go",
         "network.go",
         "profile.go",
     ],

pkg/test/dockerutil/gpu.go

@@ -0,0 +1,75 @@
+// Copyright 2023 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+// Package dockerutil provides utility functions for GPU tests.
+package dockerutil
+
+import (
+	"flag"
+
+	"github.com/docker/docker/api/types/container"
+	"github.com/docker/docker/api/types/mount"
+)
+
+// Flags.
+var (
+	setCOSGPU = flag.Bool("cos-gpu", false, "set to configure GPU settings for COS, as opposed to Docker")
+)
+
+// GPURunOpts returns Docker run options with GPU support enabled.
+func GPURunOpts() RunOpts {
+	if !*setCOSGPU {
+		return RunOpts{
+			DeviceRequests: []container.DeviceRequest{
+				{
+					Count:        -1,
+					Capabilities: [][]string{[]string{"gpu"}},
+					Options:      map[string]string{},
+				},
+			},
+		}
+	}
+
+	// COS has specific settings since it has a custom installer for GPU drivers.
+	// See: https://cloud.google.com/container-optimized-os/docs/how-to/run-gpus#install-driver
+	devices := []container.DeviceMapping{}
+	nvidia0Device := "/dev/nvidia0"
+	nvidiaUvmDevice := "/dev/nvidia-uvm"
+	nvidiactlDevice := "/dev/nvidiactl"
+	for _, device := range []string{nvidia0Device, nvidiaUvmDevice, nvidiactlDevice} {
+		devices = append(devices, container.DeviceMapping{
+			PathOnHost:        device,
+			PathInContainer:   device,
+			CgroupPermissions: "rwm",
+		})
+	}
+
+	mounts := []mount.Mount{
+		{
+			Source: "/var/lib/nvidia/lib64",
+			Target: "/usr/local/nvidia/lib64",
+			Type:   mount.TypeBind,
+		},
+		{
+			Source: "/var/lib/nvidia/bin",
+			Target: "/usr/local/nvidia/bin",
+			Type:   mount.TypeBind,
+		},
+	}
+
+	return RunOpts{
+		Mounts:  mounts,
+		Devices: devices,
+	}
+}

test/gpu/BUILD

@@ -14,9 +14,5 @@ go_test(
         "notap",
     ],
     visibility = ["//:sandbox"],
-    deps = [
-        "//pkg/test/dockerutil",
-        "@com_github_docker_docker//api/types/container:go_default_library",
-        "@com_github_docker_docker//api/types/mount:go_default_library",
-    ],
+    deps = ["//pkg/test/dockerutil"],
 )

test/gpu/gpu_smoke_test.go

@@ -12,28 +12,22 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package cos_gpu_test tests that GPUs work on Container Optimized OS (COS) images in GCP. This
-// will probably only work on COS images.
-package cos_gpu_test
+// Package gpu_smoke_test tests basic GPU functionality.
+package gpu_smoke_test
 
 import (
 	"context"
-	"flag"
 	"testing"
 
-	"github.com/docker/docker/api/types/container"
-	"github.com/docker/docker/api/types/mount"
 	"gvisor.dev/gvisor/pkg/test/dockerutil"
 )
 
-var setCOSGPU = flag.Bool("cos-gpu", false, "set to configure GPU settings for cos")
-
 func TestGPUHello(t *testing.T) {
 	ctx := context.Background()
 	c := dockerutil.MakeContainer(ctx, t)
 	defer c.CleanUp(ctx)
 
-	opts := getGPURunOpts()
+	opts := dockerutil.GPURunOpts()
 	opts.Image = "basic/cuda-vector-add"
 	out, err := c.Run(ctx, opts)
 	if err != nil {
@@ -47,57 +41,11 @@ func TestCUDATests(t *testing.T) {
 	c := dockerutil.MakeContainer(ctx, t)
 	defer c.CleanUp(ctx)
 
-	opts := getGPURunOpts()
+	opts := dockerutil.GPURunOpts()
 	opts.Image = "gpu/cuda-tests"
 	out, err := c.Run(ctx, opts)
 	if err != nil {
 		t.Fatalf("could not run cuda-tests: %v", err)
 	}
 	t.Logf("cuda-tests output: %s", string(out))
 }
-
-func getGPURunOpts() dockerutil.RunOpts {
-	if !*setCOSGPU {
-		return dockerutil.RunOpts{
-			DeviceRequests: []container.DeviceRequest{
-				{
-					Count:        -1,
-					Capabilities: [][]string{[]string{"gpu"}},
-					Options:      map[string]string{},
-				},
-			},
-		}
-	}
-
-	// COS has specific settings since it has a custom installer for GPU drivers.
-	// See: https://cloud.google.com/container-optimized-os/docs/how-to/run-gpus#install-driver
-	devices := []container.DeviceMapping{}
-	nvidia0Device := "/dev/nvidia0"
-	nvidiaUvmDevice := "/dev/nvidia-uvm"
-	nvidiactlDevice := "/dev/nvidiactl"
-	for _, device := range []string{nvidia0Device, nvidiaUvmDevice, nvidiactlDevice} {
-		devices = append(devices, container.DeviceMapping{
-			PathOnHost:        device,
-			PathInContainer:   device,
-			CgroupPermissions: "rwm",
-		})
-	}
-
-	mounts := []mount.Mount{
-		{
-			Source: "/var/lib/nvidia/lib64",
-			Target: "/usr/local/nvidia/lib64",
-			Type:   mount.TypeBind,
-		},
-		{
-			Source: "/var/lib/nvidia/bin",
-			Target: "/usr/local/nvidia/bin",
-			Type:   mount.TypeBind,
-		},
-	}
-
-	return dockerutil.RunOpts{
-		Mounts:  mounts,
-		Devices: devices,
-	}
-}