Add set_cuda_backend Context Manager to publicly expose the BETA CUDA interface

Author: NicolasHug

src/torchcodec/__init__.py

@@ -8,6 +8,7 @@
 # but that results in circular import.
 from ._frame import AudioSamples, Frame, FrameBatch  # usort:skip # noqa
 from . import decoders, samplers  # noqa
+from .decoders._decoder_utils import set_cuda_backend  # noqa
 
 try:
     # Note that version.py is generated during install.

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, Optional, 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_current_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_current_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_current_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

@@ -1702,7 +1702,7 @@ 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")