Specify all capabilities in TestGPUCheckpointRestore.

This test was broken by 5e6589e ("Update CUDA test compatibility to keep
up with added gVisor support.") which requires all
images/gpu/cuda-tests/run_sample.go users to specify "all" driver capabilities.

PiperOrigin-RevId: 713156072

pkg/test/dockerutil/gpu.go

@@ -41,13 +41,12 @@ const (
 )
 
 const (
-	// AllGPUCapabilities is the environment variable that enables all NVIDIA
+	// AllGPUCapabilitiesEnv is the environment variable that enables all NVIDIA
 	// GPU capabilities within a container.
-	AllGPUCapabilities = "NVIDIA_DRIVER_CAPABILITIES=all"
+	AllGPUCapabilitiesEnv = "NVIDIA_DRIVER_CAPABILITIES=all"
 
-	// DefaultGPUCapabilities is the environment variable that enables default
-	// NVIDIA GPU capabilities within a container.
-	DefaultGPUCapabilities = "NVIDIA_DRIVER_CAPABILITIES=compute,utility"
+	// DefaultGPUCapabilities are the driver capabilities enabled by default.
+	DefaultGPUCapabilities = "compute,utility"
 )
 
 // GPURunOpts returns Docker run options with GPU support enabled.
@@ -79,7 +78,7 @@ func GPURunOpts(sniffGPUOpts SniffGPUOpts) (RunOpts, error) {
 			ReadOnly: true,
 		})
 	}
-	gpuEnv := []string{sniffGPUOpts.GPUCapabilities()}
+	gpuEnv := []string{"NVIDIA_DRIVER_CAPABILITIES=" + sniffGPUOpts.GPUCapabilities()}
 
 	if !*setCOSGPU {
 		return RunOpts{

test/gpu/cuda_test.go

@@ -433,7 +433,7 @@ func (*FullyCompatible) IsExpectedFailure(ctx context.Context, env *TestEnvironm
 // getContainerOpts returns the container run options to run CUDA tests.
 func getContainerOpts() (dockerutil.RunOpts, error) {
 	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{
-		Capabilities: dockerutil.AllGPUCapabilities,
+		Capabilities: dockerutil.AllGPUCapabilitiesEnv,
 	})
 	if err != nil {
 		return dockerutil.RunOpts{}, fmt.Errorf("failed to get GPU run options: %w", err)
@@ -885,7 +885,7 @@ func TestCUDA(t *testing.T) {
 			t.Errorf(
 				"  $ docker run --runtime=%s --gpus=all -e %s --rm %s /run_sample %s",
 				dockerutil.Runtime(),
-				dockerutil.AllGPUCapabilities,
+				dockerutil.AllGPUCapabilitiesEnv,
 				runOpts.Image,
 				failedTests[0],
 			)

test/gpu/ffmpeg_test.go

@@ -34,7 +34,7 @@ func TestFffmpegEncodeGPU(t *testing.T) {
 	container := dockerutil.MakeContainer(ctx, t)
 	defer container.CleanUp(ctx)
 	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{
-		Capabilities:           "NVIDIA_DRIVER_CAPABILITIES=video",
+		Capabilities:           "video",
 		AllowIncompatibleIoctl: true, // TODO(gvisor.dev/issue/9452): Remove once supported in gVisor.
 	})
 	if err != nil {
@@ -60,7 +60,7 @@ func TestFffmpegDecodeGPU(t *testing.T) {
 	container := dockerutil.MakeContainer(ctx, t)
 	defer container.CleanUp(ctx)
 	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{
-		Capabilities:           "NVIDIA_DRIVER_CAPABILITIES=video",
+		Capabilities:           "video",
 		AllowIncompatibleIoctl: true, // TODO(gvisor.dev/issue/9452): Remove once supported in gVisor.
 	})
 	if err != nil {

test/gpu/sr_test.go

@@ -35,7 +35,9 @@ func TestGPUCheckpointRestore(t *testing.T) {
 	c := dockerutil.MakeContainer(ctx, t)
 	defer c.CleanUp(ctx)
 
-	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{})
+	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{
+		Capabilities: "all",
+	})
 	if err != nil {
 		t.Fatalf("failed to get GPU run options: %v", err)
 	}