Let device default be None on blocks (#1656)

Author: NicolasHug

examples/decoding/blocks.py

@@ -65,7 +65,8 @@
 #
 # ``PacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``:
 # decoding then runs on NVDEC and the color conversion on the GPU, and the
-# frames never leave the device. Demuxing always happens on the CPU.
+# frames never leave the device. Demuxing always happens on the CPU. Left
+# unspecified, ``device`` is the current default device.
 from torchcodec.decoders._blocks import ColorConverter, Demuxer, PacketDecoder
 
 demuxer = Demuxer(video_path)

src/torchcodec/decoders/_blocks/_color_converter.py

@@ -13,7 +13,7 @@
 from torchcodec._core.ops import _blocks_convert_frame, _blocks_create_color_converter
 from torchcodec._frame import Frame
 
-from .._decoder_utils import convert_output_dtype_to_str
+from .._decoder_utils import convert_device_to_str, convert_output_dtype_to_str
 from ._frame import DecodedFrame
 
 
@@ -35,16 +35,19 @@ class ColorConverter:
     Rotation is applied too, so the output matches ``VideoDecoder``'s. The angle
     is part of the frame, like its dims and colorspace, so honoring it doesn't
     bind the converter to a stream either.
+
+    ``device`` accepts a string or a ``torch.device``. It defaults to ``None``,
+    which means the current default device (see ``torch.set_default_device``).
     """
 
-    # TODO_API_BREAKDOWN UF P1: device default should be None
     def __init__(
         self,
-        device="cpu",
+        device: str | torch.device | None = None,
         output_dtype: torch.dtype | Literal["auto"] = torch.uint8,
     ):
         self._handle = _blocks_create_color_converter(
-            device=device, output_dtype=convert_output_dtype_to_str(output_dtype)
+            device=convert_device_to_str(device),
+            output_dtype=convert_output_dtype_to_str(output_dtype),
         )
 
     def convert(self, decoded_frame: DecodedFrame) -> Frame:

src/torchcodec/decoders/_blocks/_packet_decoder.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import torch
+
 from torchcodec._core.ops import (
     _blocks_create_packet_decoder,
     _blocks_packet_decoder_receive_frame,
@@ -14,6 +16,7 @@
     _blocks_packet_decoder_send_packet,
 )
 
+from .._decoder_utils import convert_device_to_str
 from ._demuxer import Demuxer
 from ._frame import DecodedFrame, Packet
 
@@ -30,12 +33,14 @@ class PacketDecoder:
     use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count
     is kept at 1 for now (not exposed); parallelism comes from composing blocks
     on your own threads.
+
+    ``device`` accepts a string or a ``torch.device``. It defaults to ``None``,
+    which means the current default device (see ``torch.set_default_device``).
     """
 
-    # TODO_API_BREAKDOWN UF P1: device default should be None, here and everywhere else
-    def __init__(self, demuxer: Demuxer, device="cpu"):
+    def __init__(self, demuxer: Demuxer, device: str | torch.device | None = None):
         self._handle = _blocks_create_packet_decoder(
-            demuxer._handle, num_threads=1, device=device
+            demuxer._handle, num_threads=1, device=convert_device_to_str(device)
         )
         self._drained = False
 

src/torchcodec/decoders/_decoder_utils.py

@@ -107,6 +107,14 @@ def get_nvdec_cache_capacity() -> int:
     return _core.get_nvdec_cache_capacity()
 
 
+def convert_device_to_str(device: str | torch.device | None) -> str:
+    # The core ops take the device as a string. None means "the current default
+    # device".
+    if device is None:
+        device = torch.get_default_device()
+    return str(device)
+
+
 def convert_output_dtype_to_str(output_dtype: torch.dtype | str) -> str:
     # The core ops take the dtype as a string.
     dtype_to_str = {torch.uint8: "uint8", torch.float32: "float32"}

src/torchcodec/decoders/_video_decoder.py

@@ -20,6 +20,7 @@
 from torchcodec._logging import _LG
 from torchcodec.decoders._decoder_utils import (
     _get_cuda_backend,
+    convert_device_to_str,
     convert_output_dtype_to_str,
 )
 from torchcodec.transforms import DecoderTransform
@@ -227,10 +228,7 @@ def __init__(
         output_dtype = convert_output_dtype_to_str(output_dtype)
 
         device_variant = _get_cuda_backend()
-        if device is None:
-            device = str(torch.get_default_device())
-        elif isinstance(device, torch_device):
-            device = str(device)
+        device = convert_device_to_str(device)
 
         if (
             device.startswith("cuda")

test/test_decoders.py

@@ -3902,6 +3902,40 @@ def _first_frame(self, path, device):
         frame = next(self._decode(decoder, self._demux(demuxer)))
         return frame, converter
 
+    @pytest.mark.parametrize("device_str", _block_devices())
+    def test_device_none_default_device(self, device_str):
+        # PacketDecoder and ColorConverter default to device=None, which should
+        # respect both the torch.device() context manager and
+        # torch.set_default_device().
+
+        def assert_first_frame_is_on_default_device():
+            # Note the absence of any device parameter.
+            demuxer = Demuxer(NASA_VIDEO.path)
+            decoder = PacketDecoder(demuxer)
+            converter = ColorConverter()
+            decoded = next(self._decode(decoder, self._demux(demuxer)))
+            # DecodedFrame.device is a string, and it carries an index
+            # ("cuda:0") since that's what torch.get_default_device() reports.
+            assert torch.device(decoded.device).type == device_str
+            assert converter.convert(decoded).data.device.type == device_str
+
+        with torch.device(device_str):
+            assert_first_frame_is_on_default_device()
+
+        original_device = torch.get_default_device()
+        try:
+            torch.set_default_device(device_str)
+            assert_first_frame_is_on_default_device()
+        finally:
+            torch.set_default_device(original_device)
+
+    @pytest.mark.parametrize("device", _block_devices())
+    def test_device_torch_device_instance(self, device):
+        # device can be a torch.device instance, not just a string.
+        frame, converter = self._first_frame(NASA_VIDEO.path, torch.device(device))
+        assert torch.device(frame.device).type == device
+        assert converter.convert(frame).data.device.type == device
+
     @pytest.mark.parametrize("case", _MATERIALIZE_VIDEOS, ids=_materialize_ids)
     @pytest.mark.parametrize("device", _block_devices())
     def test_materialize_structure(self, case, device):