From b89f53b2ceb0ddb31d93cb4f7e461e2ae4443fb5 Mon Sep 17 00:00:00 2001
From: Etienne Perot <eperot@google.com>
Date: Thu, 3 Oct 2024 19:30:36 -0700
Subject: [PATCH] Add `ffmpeg` GPU test with `h264_nvenc` video codec (which
 uses NVENC).

This test does NOT work yet in gVisor.

Updates #9452

PiperOrigin-RevId: 682127429
---
 Makefile                |  4 +++-
 test/gpu/BUILD          | 14 ++++++++++++
 test/gpu/ffmpeg_test.go | 48 +++++++++++++++++++++++++++++++++++++++++
 3 files changed, 65 insertions(+), 1 deletion(-)
 create mode 100644 test/gpu/ffmpeg_test.go

diff --git a/Makefile b/Makefile
index 8bfe91bfad..99eb20775d 100644
--- a/Makefile
+++ b/Makefile
@@ -304,7 +304,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 load-gpu_nccl-tests
+gpu-images: gpu-smoke-images load-gpu_pytorch load-gpu_ollama load-gpu_ollama_client load-basic_busybox load-basic_alpine load-basic_python load-gpu_stable-diffusion-xl load-gpu_vllm load-gpu_nccl-tests load-benchmarks_ffmpeg
 .PHONY: gpu-images
 
 gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
@@ -314,6 +314,7 @@ gpu-all-tests: gpu-images gpu-smoke-tests $(RUNTIME_BIN)
 	@$(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))
+	@$(call sudo,test/gpu:ffmpeg_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 	@$(call sudo,test/gpu:sniffer_test,--runtime=$(RUNTIME) -test.v $(ARGS))
 .PHONY: gpu-all-tests
 
@@ -324,6 +325,7 @@ cos-gpu-all-tests: gpu-images cos-gpu-smoke-tests $(RUNTIME_BIN)
 	@$(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))
+	@$(call sudo,test/gpu:ffmpeg_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 	@$(call sudo,test/gpu:sniffer_test,--runtime=$(RUNTIME) -test.v --cos-gpu $(ARGS))
 .PHONY: cos-gpu-all-tests
 
diff --git a/test/gpu/BUILD b/test/gpu/BUILD
index 25317239ea..290060f3e2 100644
--- a/test/gpu/BUILD
+++ b/test/gpu/BUILD
@@ -89,6 +89,20 @@ go_test(
     ],
 )
 
+go_test(
+    name = "ffmpeg_test",
+    srcs = ["ffmpeg_test.go"],
+    # runsc is needed to invalidate the bazel cache in case of any code changes.
+    data = ["//runsc"],
+    tags = [
+        "manual",
+        "noguitar",
+        "notap",
+    ],
+    visibility = ["//:sandbox"],
+    deps = ["//pkg/test/dockerutil"],
+)
+
 go_test(
     name = "imagegen_test",
     srcs = ["imagegen_test.go"],
diff --git a/test/gpu/ffmpeg_test.go b/test/gpu/ffmpeg_test.go
new file mode 100644
index 0000000000..3c74bf95f9
--- /dev/null
+++ b/test/gpu/ffmpeg_test.go
@@ -0,0 +1,48 @@
+// Copyright 2020 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 ffmpeg_test
+
+import (
+	"context"
+	"strings"
+	"testing"
+
+	"gvisor.dev/gvisor/pkg/test/dockerutil"
+)
+
+// TestFffmpegGPU runs ffmpeg in a GPU container using NVENC.
+func TestFffmpegGPU(t *testing.T) {
+	ctx := context.Background()
+	isGVisor, err := dockerutil.IsGVisorRuntime(ctx, t)
+	if err != nil {
+		t.Fatalf("Failed to determine if runtime is gVisor: %v", err)
+	}
+	if isGVisor {
+		t.Skip("This test is currently broken in gVisor")
+	}
+	container := dockerutil.MakeContainer(ctx, t)
+	defer container.CleanUp(ctx)
+	opts, err := dockerutil.GPURunOpts(dockerutil.SniffGPUOpts{
+		Capabilities:           "NVIDIA_DRIVER_CAPABILITIES=video",
+		AllowIncompatibleIoctl: true,
+	})
+	if err != nil {
+		t.Fatalf("Failed to get GPU run options: %v", err)
+	}
+	opts.Image = "benchmarks/ffmpeg"
+	cmd := strings.Split("ffmpeg -i video.mp4 -c:v h264_nvenc -preset fast output.mp4", " ")
+	if output, err := container.Run(ctx, opts, cmd...); err != nil {
+		t.Errorf("failed to run container: %v; output:\n%s", err, output)
+	}
+}