diff --git a/images/gpu/nccl-tests/Dockerfile.x86_64 b/images/gpu/nccl-tests/Dockerfile similarity index 100% rename from images/gpu/nccl-tests/Dockerfile.x86_64 rename to images/gpu/nccl-tests/Dockerfile diff --git a/images/gpu/ollama/Dockerfile.x86_64 b/images/gpu/ollama/Dockerfile similarity index 100% rename from images/gpu/ollama/Dockerfile.x86_64 rename to images/gpu/ollama/Dockerfile diff --git a/pkg/test/dockerutil/dockerutil.go b/pkg/test/dockerutil/dockerutil.go index 8f8e1d610f..c2dcc5579a 100644 --- a/pkg/test/dockerutil/dockerutil.go +++ b/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 diff --git a/test/gpu/BUILD b/test/gpu/BUILD index fa819b278c..d4323e357e 100644 --- a/test/gpu/BUILD +++ b/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( diff --git a/test/gpu/imagegen_test.go b/test/gpu/imagegen_test.go index ab7995bc88..9428782084 100644 --- a/test/gpu/imagegen_test.go +++ b/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) diff --git a/test/gpu/pytorch_test.go b/test/gpu/pytorch_test.go index bf9ca9f403..07acea7233 100644 --- a/test/gpu/pytorch_test.go +++ b/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{}) diff --git a/test/gpu/smoke_test.go b/test/gpu/smoke_test.go index 40ea8be3d7..81fce38920 100644 --- a/test/gpu/smoke_test.go +++ b/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) diff --git a/test/gpu/vllm/vllm_test.go b/test/gpu/vllm/vllm_test.go index 5ee98a2228..8cd33e5c21 100644 --- a/test/gpu/vllm/vllm_test.go +++ b/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) }