From a6613282eab1716edd8ecc617f2e3cb73b0929e5 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 14 Oct 2025 15:28:50 +0100 Subject: [PATCH 1/3] Add set_cuda_backend Context Manager to publicly expose the BETA CUDA interface --- src/torchcodec/__init__.py | 1 + src/torchcodec/decoders/_decoder_utils.py | 53 ++++++++++++++++++++++- src/torchcodec/decoders/_video_decoder.py | 23 +++++----- test/test_decoders.py | 2 +- 4 files changed, 66 insertions(+), 13 deletions(-) diff --git a/src/torchcodec/__init__.py b/src/torchcodec/__init__.py index 2b6c9fb3e..9525bd392 100644 --- a/src/torchcodec/__init__.py +++ b/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. diff --git a/src/torchcodec/decoders/_decoder_utils.py b/src/torchcodec/decoders/_decoder_utils.py index 3248f8362..10732b3bf 100644 --- a/src/torchcodec/decoders/_decoder_utils.py +++ b/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() diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index 729fd4727..9b7af6bb2 100644 --- a/src/torchcodec/decoders/_video_decoder.py +++ b/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, diff --git a/test/test_decoders.py b/test/test_decoders.py index 297548f16..1738336bb 100644 --- a/test/test_decoders.py +++ b/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") From 5d0eeaf32a491701985dc083385bc21dbc345f8b Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Tue, 14 Oct 2025 15:33:17 +0100 Subject: [PATCH 2/3] lint --- src/torchcodec/decoders/_decoder_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/torchcodec/decoders/_decoder_utils.py b/src/torchcodec/decoders/_decoder_utils.py index 10732b3bf..b30ee64e7 100644 --- a/src/torchcodec/decoders/_decoder_utils.py +++ b/src/torchcodec/decoders/_decoder_utils.py @@ -9,7 +9,7 @@ from contextlib import contextmanager from pathlib import Path -from typing import Generator, Optional, Union +from typing import Generator, Union from torch import Tensor from torchcodec import _core as core From 5a8e50b141999417ec5d3403bcc71d85161d65cb Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Wed, 15 Oct 2025 10:18:44 +0100 Subject: [PATCH 3/3] Add tests --- src/torchcodec/__init__.py | 1 - src/torchcodec/decoders/__init__.py | 1 + src/torchcodec/decoders/_decoder_utils.py | 2 +- src/torchcodec/decoders/_video_decoder.py | 4 +- test/test_decoders.py | 56 +++++++++++++++++++++++ 5 files changed, 60 insertions(+), 4 deletions(-) diff --git a/src/torchcodec/__init__.py b/src/torchcodec/__init__.py index 9525bd392..2b6c9fb3e 100644 --- a/src/torchcodec/__init__.py +++ b/src/torchcodec/__init__.py @@ -8,7 +8,6 @@ # 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. diff --git a/src/torchcodec/decoders/__init__.py b/src/torchcodec/decoders/__init__.py index 7b27a3bf4..980ba98a9 100644 --- a/src/torchcodec/decoders/__init__.py +++ b/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 diff --git a/src/torchcodec/decoders/_decoder_utils.py b/src/torchcodec/decoders/_decoder_utils.py index b30ee64e7..549756b81 100644 --- a/src/torchcodec/decoders/_decoder_utils.py +++ b/src/torchcodec/decoders/_decoder_utils.py @@ -99,5 +99,5 @@ def set_cuda_backend(backend: str) -> Generator[None, None, None]: _CUDA_BACKEND.reset(previous_state) -def _get_current_cuda_backend() -> str: +def _get_cuda_backend() -> str: return _CUDA_BACKEND.get() diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index 9b7af6bb2..f22f5a3fc 100644 --- a/src/torchcodec/decoders/_video_decoder.py +++ b/src/torchcodec/decoders/_video_decoder.py @@ -15,7 +15,7 @@ from torchcodec import _core as core, Frame, FrameBatch from torchcodec.decoders._decoder_utils import ( - _get_current_cuda_backend, + _get_cuda_backend, create_decoder, ERROR_REPORTING_INSTRUCTIONS, ) @@ -144,7 +144,7 @@ def __init__( if isinstance(device, torch_device): device = str(device) - device_variant = _get_current_cuda_backend() + device_variant = _get_cuda_backend() if device_variant == "ffmpeg": # TODONVDEC P2 rename 'default' into 'ffmpeg' everywhere. device_variant = "default" diff --git a/test/test_decoders.py b/test/test_decoders.py index 1738336bb..300c953bf 100644 --- a/test/test_decoders.py +++ b/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, @@ -1705,6 +1707,60 @@ def test_beta_cuda_interface_error(self): 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))