Add NCCL tests as a regression test.

PiperOrigin-RevId: 651541375

Makefile

@@ -298,7 +298,7 @@ cos-gpu-smoke-tests: gpu-smoke-images $(RUNTIME_BIN)
 # This is a superset of those needed for smoke tests.
 # It includes non-GPU images that are used as part of GPU tests,
 # e.g. busybox and python.
-gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python load-gpu_stable-diffusion-xl load-gpu_vllm
+gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_python load-gpu_stable-diffusion-xl load-gpu_vllm load-gpu_nccl-tests
 .PHONY: gpu-images
 
 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
@@ -307,6 +307,7 @@ gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v $(ARGS))
+	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
 cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
@@ -315,6 +316,7 @@ cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(call sudo,test/gpu:textgen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:imagegen_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sr_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
+	@$(call sudo,test/gpu:nccl_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests
 
 portforward-tests: load-basic_redis load-basic_nginx $(RUNTIME_BIN)

images/gpu/nccl-tests/Dockerfile.x86_64

@@ -0,0 +1,6 @@
+FROM nvidia/cuda:12.5.0-devel-ubuntu22.04
+
+RUN apt-get update && apt-get install git -y
+
+RUN git clone --depth=1 https://github.com/NVIDIA/nccl-tests.git && \
+    cd nccl-tests && make

test/gpu/BUILD

@@ -121,3 +121,15 @@ go_test(
         "@com_github_docker_docker//api/types/mount:go_default_library",
     ],
 )
+
+go_test(
+    name = "nccl_test",
+    srcs = ["nccl_test.go"],
+    tags = [
+        "manual",
+        "noguitar",
+        "notap",
+    ],
+    visibility = ["//:sandbox"],
+    deps = ["//pkg/test/dockerutil"],
+)

test/gpu/nccl_test.go

@@ -0,0 +1,60 @@
+// Copyright 2024 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 nccl_test runs through NCCL tests.
+package nccl_test
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+)
+
+// runNCCL runs the given script and command in a NCCL container.
+func runNCCL(ctx context.Context, t *testing.T, testName string) {
+	t.Helper()
+	c := dockerutil.MakeContainer(ctx, t)
+	opts := dockerutil.GPURunOpts()
+	opts.Image = "gpu/nccl-tests"
+	cmd := fmt.Sprintf("/nccl-tests/build/%s", testName)
+	out, err := c.Run(ctx, opts, cmd)
+	if err != nil {
+		t.Errorf("Failed: %v\nContainer output:\n%s", err, out)
+	} else {
+		t.Logf("Container output:\n%s", out)
+	}
+}
+
+func TestNCCL(t *testing.T) {
+	testNames := []string{
+		"all_gather_perf",
+		"all_reduce_perf",
+		"alltoall_perf",
+		"broadcast_perf",
+		"gather_perf",
+		"hypercube_perf",
+		"reduce_perf",
+		"reduce_scatter_perf",
+		"scatter_perf",
+		"sendrecv_perf",
+	}
+
+	ctx := context.Background()
+	for _, test := range testNames {
+		t.Logf("Running NCCL test: %s", test)
+		runNCCL(ctx, t, test)
+	}
+}