Enable CUDA support with FFmpeg4 (#517)

Author: NicolasHug

.github/workflows/linux_cuda_wheel.yaml

@@ -68,7 +68,11 @@ jobs:
         python-version: ['3.9']
         cuda-version: ['11.8', '12.4', '12.6']
         # TODO: put back ffmpeg 5 https://github.com/pytorch/torchcodec/issues/325
-        ffmpeg-version-for-tests: ['6', '7']
+        ffmpeg-version-for-tests: ['4.4.2', '6', '7']
+        exclude:
+        - cuda-version: '12.4'  # TODO put this back it fails with infra issue.
+          ffmpeg-version-for-tests: '4.4.2'
+
     container:
       image: "pytorch/manylinux2_28-builder:cuda${{ matrix.cuda-version }}"
       options: "--gpus all -e NVIDIA_DRIVER_CAPABILITIES=video,compute,utility"
@@ -141,7 +145,7 @@ jobs:
           ${CONDA_RUN} python test/decoders/manual_smoke_test.py
       - name: Run Python tests
         run: |
-          ${CONDA_RUN} FAIL_WITHOUT_CUDA=1 pytest test -vvv
+          ${CONDA_RUN} FAIL_WITHOUT_CUDA=1 pytest test -v --tb=short
       - name: Run Python benchmark
         run: |
           ${CONDA_RUN} time python benchmarks/decoders/gpu_benchmark.py --devices=cuda:0,cpu --resize_devices=none

src/torchcodec/decoders/_core/CMakeLists.txt

@@ -77,10 +77,7 @@ if(DEFINED ENV{BUILD_AGAINST_ALL_FFMPEG_FROM_S3})
     )
 
 
-    if(NOT ENABLE_CUDA)
-	    # TODO: Enable more ffmpeg versions for cuda.
-	    make_torchcodec_library(libtorchcodec4 ffmpeg4)
-	endif()
+	make_torchcodec_library(libtorchcodec4 ffmpeg4)
 	make_torchcodec_library(libtorchcodec7 ffmpeg7)
 	make_torchcodec_library(libtorchcodec6 ffmpeg6)
 	make_torchcodec_library(libtorchcodec5 ffmpeg5)

test/decoders/test_video_decoder.py

@@ -15,8 +15,8 @@
     assert_frames_equal,
     AV1_VIDEO,
     cpu_and_cuda,
+    get_ffmpeg_major_version,
     H265_VIDEO,
-    in_fbcode,
     NASA_VIDEO,
 )
 
@@ -249,7 +249,7 @@ def test_getitem_slice(self, device, seek_mode):
             ]
         )
         for sliced, ref in zip(all_frames, decoder):
-            if not (in_fbcode() and device == "cuda"):
+            if not (device == "cuda" and get_ffmpeg_major_version() == 4):
                 # TODO: remove the "if".
                 # See https://github.com/pytorch/torchcodec/issues/428
                 assert_frames_equal(sliced, ref)
@@ -428,7 +428,7 @@ def test_get_frames_at_fails(self, device, seek_mode):
 
     @pytest.mark.parametrize("device", cpu_and_cuda())
     def test_get_frame_at_av1(self, device):
-        if in_fbcode() and device == "cuda":
+        if device == "cuda" and get_ffmpeg_major_version() == 4:
             return
 
         decoder = VideoDecoder(AV1_VIDEO.path, device=device)

test/utils.py

@@ -11,6 +11,8 @@
 
 import torch
 
+from torchcodec.decoders._core import get_ffmpeg_library_versions
+
 
 # Decorator for skipping CUDA tests when CUDA isn't available. The tests are
 # effectively marked to be skipped in pytest_collection_modifyitems() of
@@ -23,6 +25,10 @@ def cpu_and_cuda():
     return ("cpu", pytest.param("cuda", marks=pytest.mark.needs_cuda))
 
 
+def get_ffmpeg_major_version():
+    return int(get_ffmpeg_library_versions()["ffmpeg_version"].split(".")[0])
+
+
 # For use with decoded data frames. On CPU Linux, we expect exact, bit-for-bit
 # equality. On CUDA Linux, we expect a small tolerance.
 # On other platforms (e.g. MacOS), we also allow a small tolerance. FFmpeg does
@@ -32,7 +38,7 @@ def assert_frames_equal(*args, **kwargs):
     if sys.platform == "linux":
         if args[0].device.type == "cuda":
             atol = 2
-            if in_fbcode():
+            if get_ffmpeg_major_version() == 4:
                 assert_tensor_close_on_at_least(
                     args[0], args[1], percentage=95, atol=atol
                 )