Update GPU tests to work on ARM

pkg/test/dockerutil/dockerutil.go

@@ -25,6 +25,7 @@ import (
 	"os"
 	"os/exec"
 	"regexp"
+	rt "runtime"
 	"strconv"
 	"strings"
 	"testing"
@@ -242,6 +243,12 @@ func runtimeMap() (map[string]any, error) {
 	return rs, nil
 }
 
+// IsRunningOnARM returns true if the current runtime is running on an ARM
+// machine.
+func IsRunningOnARM() bool {
+	return strings.HasPrefix(rt.GOARCH, "arm")
+}
+
 // Save exports a container image to the given Writer.
 //
 // Note that the writer should be actively consuming the output, otherwise it

test/gpu/BUILD

@@ -114,7 +114,10 @@ go_test(
         "notap",
     ],
     visibility = ["//:sandbox"],
-    deps = ["//test/gpu/stablediffusion"],
+    deps = [
+        "//pkg/test/dockerutil",
+        "//test/gpu/stablediffusion",
+    ],
 )
 
 go_test(

test/gpu/imagegen_test.go

@@ -20,11 +20,16 @@ import (
 	"testing"
 	"time"
 
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
 	"gvisor.dev/gvisor/test/gpu/stablediffusion"
 )
 
 // TestStableDiffusionXL generates an image with Stable Diffusion XL.
 func TestStableDiffusionXL(t *testing.T) {
+	if dockerutil.IsRunningOnARM() {
+		fbIssue := "https://github.com/facebookresearch/xformers/issues/1071"
+		t.Skipf("%s does not run on ARM due to library issues: see: %q", t.Name(), fbIssue)
+	}
 	ctx := context.Background()
 	sdxl := stablediffusion.NewDockerXL(t)
 	generateCtx, generateCancel := context.WithTimeout(ctx, 15*time.Minute)

test/gpu/pytorch_test.go

@@ -24,6 +24,9 @@ import (
 
 // runPytorch runs the given script and command in a PyTorch container.
 func runPytorch(ctx context.Context, t *testing.T, scriptPath string, args ...string) {
+	if dockerutil.IsRunningOnARM() {
+		t.Skipf("%s is not supported on ARM for now.", t.Name())
+	}
 	t.Helper()
 	c := dockerutil.MakeContainer(ctx, t)
 	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{})

test/gpu/smoke_test.go

@@ -23,6 +23,9 @@ import (
 )
 
 func TestGPUHello(t *testing.T) {
+	if dockerutil.IsRunningOnARM() {
+		t.Skipf("%s is not supported on ARM due to cross compile errors.", t.Name())
+	}
 	ctx := context.Background()
 	c := dockerutil.MakeContainer(ctx, t)
 	defer c.CleanUp(ctx)

test/gpu/vllm/vllm_test.go

@@ -33,6 +33,9 @@ import (
 
 // BenchmarkVLLM runs a vLLM workload.
 func BenchmarkVLLM(b *testing.B) {
+	if dockerutil.IsRunningOnARM() {
+		b.Skipf("%s is not supported on ARM for now.", b.Name())
+	}
 	doVLLMTest(b)
 }