Add COS GPU Pipeline.

Make updates to the gpu tests and create a pipeline for COS images.
COS has a different installation procedure for GPU drivers, and due
to it having specific file system permissions set, we need to change
how we import CUDA libraries and the GPU devices themselves.

Update the tests to do this correctly and add a COS pipeline.

PiperOrigin-RevId: 580645686

.buildkite/pipeline.yaml

@@ -174,9 +174,16 @@ steps:
     label: ":firefighter: GPU Smoke Test"
     commands:
       - make sudo TARGETS=//tools/gpu:main ARGS="install --latest" || cat /var/log/nvidia-installer.log
-      - make gpu-tests
+      - make gpu-tests RUNTIME_ARGS="--strace --debug"
     agents:
       queue: gpu
+  - <<: *common
+    <<: *source_test
+    label: ":female_supervillain: COS GPU Smoke Test"
+    commands:
+      - make cos-gpu-tests RUNTIME_ARGS="--strace --debug"
+    agents:
+      queue: cos-canary-gpu
   - <<: *common
     <<: *source_test_continuous
     label: ":screwdriver: All GPU Drivers Test"

Makefile

@@ -265,8 +265,14 @@ simple-tests: unit-tests # Compatibility target.
 gpu-tests: load-basic_cuda-vector-add load-gpu_cuda-tests $(RUNTIME_BIN)
 	@$(call test,--test_env=RUNTIME=runc //test/gpu:gpu_test)
 	@$(call install_runtime,$(RUNTIME),--platform=systrap --nvproxy=true --nvproxy-docker=true)
-	@$(call test_runtime,$(RUNTIME),//test/gpu:gpu_test)
-.PHONE: gpu-tests
+	@$(call sudo,test/gpu:gpu_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+.PHONY: gpu-tests
+
+cos-gpu-tests: load-basic_cuda-vector-add load-gpu_cuda-tests $(RUNTIME_BIN)
+	@$(call sudo,test/gpu:gpu_test,--runtime=runc -test.v --cos-gpu $(ARGS))
+	@$(call install_runtime,$(RUNTIME),--platform=systrap --nvproxy=true)
+	@$(call sudo,test/gpu:gpu_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+.PHONY: cos-gpu-tests
 
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--network=sandbox)

pkg/test/dockerutil/container.go

@@ -106,8 +106,10 @@ type RunOpts struct {
 	// Links is the list of containers to be connected to the container.
 	Links []string
 
-	// Devices are device requests on the container itself.
-	Devices []container.DeviceRequest
+	// DeviceRequests are device requests on the container itself.
+	DeviceRequests []container.DeviceRequest
+
+	Devices []container.DeviceMapping
 }
 
 func makeContainer(ctx context.Context, logger testutil.Logger, runtime string) *Container {
@@ -281,7 +283,8 @@ func (c *Container) hostConfig(r RunOpts) *container.HostConfig {
 		Resources: container.Resources{
 			Memory:         int64(r.Memory), // In bytes.
 			CpusetCpus:     r.CpusetCpus,
-			DeviceRequests: r.Devices,
+			DeviceRequests: r.DeviceRequests,
+			Devices:        r.Devices,
 		},
 	}
 }

test/gpu/BUILD

@@ -17,5 +17,6 @@ go_test(
     deps = [
         "//pkg/test/dockerutil",
         "@com_github_docker_docker//api/types/container:go_default_library",
+        "@com_github_docker_docker//api/types/mount:go_default_library",
     ],
 )

test/gpu/gpu_test.go

@@ -18,32 +18,27 @@ package cos_gpu_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)
 
-	out, err := c.Run(ctx, dockerutil.RunOpts{
-		Image: "basic/cuda-vector-add",
-		Devices: []container.DeviceRequest{
-			{
-				Count:        -1,
-				Capabilities: [][]string{[]string{"gpu"}},
-				Options:      map[string]string{},
-			},
-		},
-	})
-
+	opts := getGPURunOpts()
+	opts.Image = "basic/cuda-vector-add"
+	out, err := c.Run(ctx, opts)
 	if err != nil {
 		t.Fatalf("could not run cuda-vector-add: %v", err)
 	}
-
 	t.Logf("cuda-vector-add output: %s", string(out))
 }
 
@@ -52,20 +47,57 @@ func TestCUDATests(t *testing.T) {
 	c := dockerutil.MakeContainer(ctx, t)
 	defer c.CleanUp(ctx)
 
-	out, err := c.Run(ctx, dockerutil.RunOpts{
-		Image: "gpu/cuda-tests",
-		Devices: []container.DeviceRequest{
-			{
-				Count:        -1,
-				Capabilities: [][]string{[]string{"gpu"}},
-				Options:      map[string]string{},
-			},
-		},
-	})
-
+	opts := getGPURunOpts()
+	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,
+	}
+}