diff --git a/pkg/test/dockerutil/BUILD b/pkg/test/dockerutil/BUILD index 46cb798fd7..afa93f506d 100644 --- a/pkg/test/dockerutil/BUILD +++ b/pkg/test/dockerutil/BUILD @@ -12,6 +12,7 @@ go_library( "container.go", "dockerutil.go", "exec.go", + "gpu.go", "network.go", "profile.go", ], diff --git a/pkg/test/dockerutil/gpu.go b/pkg/test/dockerutil/gpu.go new file mode 100644 index 0000000000..a116cec86b --- /dev/null +++ b/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, + } +} diff --git a/test/gpu/BUILD b/test/gpu/BUILD index f8c5df9e4c..34056fd6aa 100644 --- a/test/gpu/BUILD +++ b/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"], ) diff --git a/test/gpu/gpu_smoke_test.go b/test/gpu/gpu_smoke_test.go index 4f172e3b36..d2576fe412 100644 --- a/test/gpu/gpu_smoke_test.go +++ b/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,7 +41,7 @@ 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 { @@ -55,49 +49,3 @@ func TestCUDATests(t *testing.T) { } 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, - } -}