Add set_cuda_backend Context Manager to publicly expose the BETA CUDA Interface

src/torchcodec/decoders/__init__.py

@@ -6,6 +6,7 @@
 
 from .._core import AudioStreamMetadata, VideoStreamMetadata
 from ._audio_decoder import AudioDecoder  # noqa
+from ._decoder_utils import set_cuda_backend  # noqa
 from ._video_decoder import VideoDecoder  # noqa
 
 SimpleVideoDecoder = VideoDecoder

src/torchcodec/decoders/_decoder_utils.py

@@ -4,10 +4,12 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import contextvars
 import io
+from contextlib import contextmanager
 from pathlib import Path
 
-from typing import Union
+from typing import Generator, Union
 
 from torch import Tensor
 from torchcodec import _core as core
@@ -50,3 +52,52 @@ def create_decoder(
         "read(self, size: int) -> bytes and "
         "seek(self, offset: int, whence: int) -> int methods."
     )
+
+
+# Thread-local and async-safe storage for the current CUDA backend
+_CUDA_BACKEND: contextvars.ContextVar[str] = contextvars.ContextVar(
+    "_CUDA_BACKEND", default="ffmpeg"
+)
+
+
+@contextmanager
+def set_cuda_backend(backend: str) -> Generator[None, None, None]:
+    """Context Manager to set the CUDA backend for :class:`~torchcodec.decoders.VideoDecoder`.
+
+    This context manager allows you to specify which CUDA backend implementation
+    to use when creating :class:`~torchcodec.decoders.VideoDecoder` instances
+    with CUDA devices.  This is thread-safe and async-safe.
+
+    Note that you still need to pass `device="cuda"` when creating the
+    :class:`~torchcodec.decoders.VideoDecoder` instance. If a CUDA device isn't
+    specified, this context manager will have no effect.
+
+    Only the creation of the decoder needs to be inside the context manager, the
+    decoding methods can be called outside of it.
+
+    Args:
+        backend (str): The CUDA backend to use. Can be "ffmpeg" or "beta". Default is "ffmpeg".
+
+    Example:
+        >>> with torchcodec.set_cuda_backend("beta"):
+        ...     decoder = VideoDecoder("video.mp4", device="cuda")
+        ...
+        ... # Only the decoder creation needs to be part of the context manager.
+        ... # Decoder will now the beta CUDA implementation:
+        ... decoder.get_frame_at(0)
+    """
+    backend = backend.lower()
+    if backend not in ("ffmpeg", "beta"):
+        raise ValueError(
+            f"Invalid CUDA backend ({backend}). Supported values are 'ffmpeg' and 'beta'."
+        )
+
+    previous_state = _CUDA_BACKEND.set(backend)
+    try:
+        yield
+    finally:
+        _CUDA_BACKEND.reset(previous_state)
+
+
+def _get_cuda_backend() -> str:
+    return _CUDA_BACKEND.get()

src/torchcodec/decoders/_video_decoder.py

@@ -15,6 +15,7 @@
 
 from torchcodec import _core as core, Frame, FrameBatch
 from torchcodec.decoders._decoder_utils import (
+    _get_cuda_backend,
     create_decoder,
     ERROR_REPORTING_INSTRUCTIONS,
 )
@@ -143,17 +144,17 @@ def __init__(
         if isinstance(device, torch_device):
             device = str(device)
 
-        # If device looks like "cuda:0:beta", make it "cuda:0" and set
-        # device_variant to "beta"
-        # TODONVDEC P2 Consider alternative ways of exposing custom device
-        # variants, and if we want this new decoder backend to be a "device
-        # variant" at all.
-        device_variant = "default"
-        if device is not None:
-            device_split = device.split(":")
-            if len(device_split) == 3:
-                device_variant = device_split[2]
-                device = ":".join(device_split[0:2])
+        device_variant = _get_cuda_backend()
+        if device_variant == "ffmpeg":
+            # TODONVDEC P2 rename 'default' into 'ffmpeg' everywhere.
+            device_variant = "default"
+
+        # Legacy support for device="cuda:0:beta" syntax
+        # TODONVDEC P2: remove support for this everywhere. This will require
+        # updating our tests.
+        if device == "cuda:0:beta":
+            device = "cuda:0"
+            device_variant = "beta"
 
         core.add_video_stream(
             self._decoder,

test/test_decoders.py

@@ -18,9 +18,11 @@
 from torchcodec.decoders import (
     AudioDecoder,
     AudioStreamMetadata,
+    set_cuda_backend,
     VideoDecoder,
     VideoStreamMetadata,
 )
+from torchcodec.decoders._decoder_utils import _get_cuda_backend
 
 from .utils import (
     all_supported_devices,
@@ -1702,9 +1704,63 @@ def test_beta_cuda_interface_small_h265(self):
 
     @needs_cuda
     def test_beta_cuda_interface_error(self):
-        with pytest.raises(RuntimeError, match="Unsupported device"):
+        with pytest.raises(RuntimeError, match="Invalid device string"):
             VideoDecoder(NASA_VIDEO.path, device="cuda:0:bad_variant")
 
+    @needs_cuda
+    def test_set_cuda_backend(self):
+        # Tests for the set_cuda_backend() context manager.
+
+        with pytest.raises(ValueError, match="Invalid CUDA backend"):
+            with set_cuda_backend("bad_backend"):
+                pass
+
+        # set_cuda_backend() is meant to be used as a context manager. Using it
+        # as a global call does nothing because the "context" is exited right
+        # away. This is a good thing, we prefer users to use it as a CM only.
+        set_cuda_backend("beta")
+        assert _get_cuda_backend() == "ffmpeg"  # Not changed to "beta".
+
+        # Case insensitive
+        with set_cuda_backend("BETA"):
+            assert _get_cuda_backend() == "beta"
+
+        def assert_decoder_uses(decoder, *, expected_backend):
+            # Assert that a decoder instance is using a given backend.
+            #
+            # We know H265_VIDEO fails on the BETA backend while it works on the
+            # ffmpeg one.
+            if expected_backend == "ffmpeg":
+                decoder.get_frame_at(0)  # this would fail if this was BETA
+            else:
+                with pytest.raises(RuntimeError, match="Video is too small"):
+                    decoder.get_frame_at(0)
+
+        # Check that the default is the ffmpeg backend
+        assert _get_cuda_backend() == "ffmpeg"
+        dec = VideoDecoder(H265_VIDEO.path, device="cuda")
+        assert_decoder_uses(dec, expected_backend="ffmpeg")
+
+        # Check the setting "beta" effectively uses the BETA backend.
+        # We also show that the affects decoder creation only. When the decoder
+        # is created with a given backend, it stays in this backend for the rest
+        # of its life. This is normal and intended.
+        with set_cuda_backend("beta"):
+            dec = VideoDecoder(H265_VIDEO.path, device="cuda")
+        assert _get_cuda_backend() == "ffmpeg"
+        assert_decoder_uses(dec, expected_backend="beta")
+        with set_cuda_backend("ffmpeg"):
+            assert_decoder_uses(dec, expected_backend="beta")
+
+        # Hacky way to ensure passing "cuda:1" is supported by both backends. We
+        # just check that there's an error when passing cuda:N where N is too
+        # high.
+        bad_device_number = torch.cuda.device_count() + 1
+        for backend in ("ffmpeg", "beta"):
+            with pytest.raises(RuntimeError, match="invalid device ordinal"):
+                with set_cuda_backend(backend):
+                    VideoDecoder(H265_VIDEO.path, device=f"cuda:{bad_device_number}")
+
 
 class TestAudioDecoder:
     @pytest.mark.parametrize("asset", (NASA_AUDIO, NASA_AUDIO_MP3, SINE_MONO_S32))