Add GPU checkpoint/restore test.

Note that GPU state is not restored. This tests that the sandbox is restored
and the GPUs are accessible and functional after restore.

Updates #9363, #9649, #9767

PiperOrigin-RevId: 590679227

Makefile

@@ -278,15 +278,15 @@ gpu-smoke-images: load-basic_cuda-vector-add load-gpu_cuda-tests
 .PHONY: gpu-smoke-images
 
 gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
-	@$(call test,--test_env=RUNTIME=runc //test/gpu:gpu_smoke_test)
+	@$(call sudo,test/gpu:smoke_test,--runtime=runc -test.v $(ARGS))
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
-	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+	@$(call sudo,test/gpu:smoke_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-smoke-tests
 
 cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
-	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=runc -test.v --cos-gpu $(ARGS))
+	@$(call sudo,test/gpu:smoke_test,--runtime=runc -test.v --cos-gpu $(ARGS))
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
-	@$(call sudo,test/gpu:gpu_smoke_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+	@$(call sudo,test/gpu:smoke_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-smoke-tests
 
 # Images needed for GPU tests.
@@ -299,11 +299,13 @@ gpu-images: gpu-smoke-images load-gpu_ollama load-basic_busybox load-basic_pytho
 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true --nvproxy-docker=true)
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call install_runtime,$(RUNTIME),--nvproxy=true)
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests
 
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)

test/gpu/BUILD

@@ -6,8 +6,8 @@ package(
 )
 
 go_test(
-    name = "gpu_smoke_test",
-    srcs = ["gpu_smoke_test.go"],
+    name = "smoke_test",
+    srcs = ["smoke_test.go"],
     tags = [
         "local",
         "noguitar",
@@ -32,3 +32,19 @@ go_test(
         "//test/gpu/ollama",
     ],
 )
+
+go_test(
+    name = "sr_test",
+    srcs = ["sr_test.go"],
+    tags = [
+        "local",
+        "noguitar",
+        "notap",
+    ],
+    visibility = ["//:sandbox"],
+    deps = [
+        "//pkg/context",
+        "//pkg/test/dockerutil",
+        "//pkg/test/testutil",
+    ],
+)

test/gpu/smoke_test.go

@@ -12,8 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package gpu_smoke_test tests basic GPU functionality.
-package gpu_smoke_test
+// Package smoke_test tests basic GPU functionality.
+package smoke_test
 
 import (
 	"context"

test/gpu/sr_test.go

@@ -0,0 +1,67 @@
+// 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 sr_test runs checkpoint/restore tests for nvproxy.
+package sr_test
+
+import (
+	"testing"
+	"time"
+
+	"gvisor.dev/gvisor/pkg/context"
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+	"gvisor.dev/gvisor/pkg/test/testutil"
+)
+
+func TestGPUCheckpointRestore(t *testing.T) {
+	if !testutil.IsCheckpointSupported() {
+		t.Skip("Checkpoint is not supported.")
+	}
+	dockerutil.EnsureDockerExperimentalEnabled()
+
+	ctx := context.Background()
+	c := dockerutil.MakeContainer(ctx, t)
+	defer c.CleanUp(ctx)
+
+	opts := dockerutil.GPURunOpts()
+	opts.Image = "basic/cuda-vector-add"
+	if err := c.Spawn(ctx, opts, "sleep", "infinity"); err != nil {
+		t.Fatalf("could not run cuda-vector-add: %v", err)
+	}
+
+	// Run the vector add program.
+	vectorAddCmd := []string{"/bin/sh", "-c", "./vectorAdd"}
+	if _, err := c.Exec(ctx, dockerutil.ExecOpts{}, vectorAddCmd...); err != nil {
+		t.Fatalf("docker exec failed: %v", err)
+	}
+
+	// Create a snapshot.
+	if err := c.Checkpoint(ctx, "test"); err != nil {
+		t.Fatalf("docker checkpoint failed: %v", err)
+	}
+	if err := c.WaitTimeout(ctx, time.Minute); err != nil {
+		t.Fatalf("wait failed: %v", err)
+	}
+
+	// Restore the snapshot.
+	// TODO(b/143498576): Remove Poll after github.com/moby/moby/issues/38963 is fixed.
+	if err := testutil.Poll(func() error { return c.Restore(ctx, "test") }, time.Minute); err != nil {
+		t.Fatalf("docker restore failed: %v", err)
+	}
+
+	// Run the vector add program again to ensure GPUs are functional.
+	if _, err := c.Exec(ctx, dockerutil.ExecOpts{}, vectorAddCmd...); err != nil {
+		t.Fatalf("docker exec failed: %v", err)
+	}
+}