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 3248f8362..549756b81 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, 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() diff --git a/src/torchcodec/decoders/_video_decoder.py b/src/torchcodec/decoders/_video_decoder.py index 729fd4727..f22f5a3fc 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_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, diff --git a/test/test_decoders.py b/test/test_decoders.py index 297548f16..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, @@ -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))