diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index ae7ca1910a9f..9fba86ce96d9 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -259,6 +259,11 @@ Each open decoder uses memory and a file handle, so usage grows with cameras and `read.video.decode-parallelism` limits concurrent video decodes across all cameras in each Dataset worker (default: 16). Set it to 1 for serial decoding. Parallel decoding applies to video windows without `image_transforms` and uses more CPU and memory. +`read.video.decoders-per-payload` (default: 4) also splits single-video windows across +independent cached decoders, within those limits. Set it to 1 to disable this splitting; +the reader automatically reduces the decoder count when overlapping windows would add +more than 50% duplicate frame requests. Additional decoders can increase initialization +cost and repeated reads. Subclass `PaimonDatasetReader` for a custom logical frame layout: diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 77f40584ee5f..6a517766eb42 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -1181,6 +1181,15 @@ class CoreOptions: "PyPaimon LeRobot Dataset worker.") ) + READ_VIDEO_DECODERS_PER_PAYLOAD: ConfigOption[int] = ( + ConfigOptions.key("read.video.decoders-per-payload") + .int_type() + .default_value(4) + .with_description( + "Maximum independent decoders for single-payload LeRobot video windows, " + "bounded by read.video.decode-parallelism and read.video.max-open-decoders.") + ) + PARQUET_COLUMN_INDEX_ENABLED: ConfigOption[bool] = ( ConfigOptions.key("parquet.filter.columnindex.enabled") .boolean_type() @@ -1951,6 +1960,9 @@ def read_video_max_open_decoders(self) -> int: def read_video_decode_parallelism(self) -> int: return self._positive_video_option(CoreOptions.READ_VIDEO_DECODE_PARALLELISM) + def read_video_decoders_per_payload(self) -> int: + return self._positive_video_option(CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD) + def _positive_video_option(self, option) -> int: raw = self.options.to_map().get(option.key()) if isinstance(raw, (bool, float)): diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 21d4bb8b4850..1ce3998ae47a 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -279,6 +279,11 @@ def _set_frame_rows( if isinstance(table_options, CoreOptions) else CoreOptions.READ_VIDEO_DECODE_PARALLELISM.default_value() ) + self._video_decoders_per_payload = ( + table_options.read_video_decoders_per_payload() + if isinstance(table_options, CoreOptions) + else CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD.default_value() + ) self._video_decode_pool = _VideoDecodePool() self._video_collators = [ VideoFrameCollator( @@ -455,6 +460,9 @@ def _get_items(self, indices, return_batch, share_memory): plans, rows, getattr(self, "_video_collators", ()), self._features, self.return_uint8, return_batch=return_batch, share_memory=share_memory, + decoders_per_payload=getattr( + self, "_video_decoders_per_payload", + CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD.default_value()), decode_pool=getattr(self, "_video_decode_pool", None), decode_parallelism=getattr( self, "_video_decode_parallelism", @@ -1439,21 +1447,56 @@ def _attach_task_labels(rows, task_names, subtask_names): def _decode_video_windows( plans, rows, collators, features, return_uint8, *, - return_batch=False, share_memory=False, decode_parallelism=1, decode_pool=None): + return_batch=False, share_memory=False, decode_parallelism=1, decode_pool=None, + decoders_per_payload=1): pool = decode_pool if decode_pool is not None else _VideoDecodePool() try: return _decode_video_windows_with_pool( plans, rows, collators, features, return_uint8, return_batch=return_batch, share_memory=share_memory, - decode_parallelism=decode_parallelism, executor=pool.get(decode_parallelism)) + decode_parallelism=decode_parallelism, executor=pool.get(decode_parallelism), + decoders_per_payload=decoders_per_payload) finally: if decode_pool is None: pool.close() +def _partition_video_windows(windows, count): + """Split windows while bounding duplicate frame requests across decoders.""" + grouped = {} + for window in windows: + # Windows with the same requested frames should always share a decoder, + # even when their output order differs. + grouped.setdefault(tuple(sorted(set(window[1]))), []).append(window) + ordered = sorted(grouped.values(), key=lambda group: min(group[0][1])) + count = min(count, len(ordered)) + serial_frames = len({ + index for _, window in windows for index in window + }) + while count > 1: + partitions = [ + [ + window + for group in ordered[len(ordered) * replica // count: + len(ordered) * (replica + 1) // count] + for window in group + ] + for replica in range(count) + ] + parallel_frames = sum(len({ + index for _, window in partition for index in window + }) for partition in partitions) + # Bound additional frame requests to 50%. This retains parallelism for + # mostly independent windows without multiplying highly overlapping IO. + if 2 * parallel_frames <= 3 * serial_frames: + return partitions + count -= 1 + return [windows] + + def _decode_video_windows_with_pool( plans, rows, collators, features, return_uint8, *, - return_batch, share_memory, decode_parallelism, executor): + return_batch, share_memory, decode_parallelism, executor, decoders_per_payload): import torch tasks = [c for c in collators if c.video_column in plans[0]["windows"]] @@ -1487,6 +1530,18 @@ def decode(collator): requests.setdefault((payload, keyframe_index), []).append(( offset, [d.frame_index for d in window])) + if len(requests) == 1 and executor is not None: + request_key, windows = next(iter(requests.items())) + count = min(decoders_per_payload, decode_parallelism, + collator.max_open_videos, len(windows)) + if count > 1: + partitions = _partition_video_windows(windows, count) + if len(partitions) > 1: + requests = { + request_key + (replica,): partition + for replica, partition in enumerate(partitions) + } + def read_frames(decoder, windows): get_frames = getattr(decoder, "get_frames_at", None) if not callable(get_frames): diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index 3d9eb812c442..a45b01b0a0cf 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -366,8 +366,8 @@ def _prepare_row(self, row): ) return output, descriptor - def _decoder(self, descriptor, keyframe_index_descriptor=None): - cache_key = (descriptor, keyframe_index_descriptor) + def _decoder(self, descriptor, keyframe_index_descriptor=None, replica=0): + cache_key = (descriptor, keyframe_index_descriptor, replica) resource = self._decoders.pop(cache_key, None) if resource is not None: self._decoders[cache_key] = resource diff --git a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py index 45a27930d17e..17fa76a12446 100644 --- a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py +++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py @@ -442,8 +442,171 @@ def blob(file_io, uri, offset, length): finally: collator.close() self.assertTrue(torch.equal(*outputs)) + outputs = [] + single_plans = [{"windows": {"camera": indices}} for indices in ( + [6, 1, 6], [0, 1, 2], [2, 3, 4], [4, 5, 7])] + for replicas in (1, 4): + collator = VideoFrameCollator( + SimpleNamespace(file_io="test"), video_column="camera", + decoder_factory=partial(_open_video_decoder, backend="pyav"), + decode_fn=identity, max_open_videos=4) + try: + with patch("pypaimon.multimodal.video.Blob.from_file", blob), \ + patch("torch.get_num_threads", return_value=1): + outputs.append(_decode_video_windows( + single_plans, rows, [collator], + {"camera": {"dtype": "video", "shape": [16, 16, 3]}}, True, + return_batch=True, decode_parallelism=4, + decoders_per_payload=replicas)["camera"]) + self.assertEqual(len(collator._decoders), replicas) + finally: + collator.close() + self.assertTrue(torch.equal(*outputs)) + + def test_single_payload_windows_reuse_independent_decoders(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + from pypaimon.multimodal.lerobot.dataset import _decode_video_windows + from pypaimon.table.row.blob import VideoFrameDescriptor + + windows = [[7, 2, 7], [0, 1, 2], [5, 6, 7], [2, 3, 4]] + rows = {i: {"camera": VideoFrameDescriptor( + "one-video", 0, 1, i, -1, 0).serialize()} for i in range(8)} + plans = [{"windows": {"camera": window}} for window in windows] + features = {"camera": {"dtype": "video", "shape": [4, 5, 3]}} + for parallelism, capacity, replicas, expected in ( + (1, 4, 4, 1), (4, 1, 4, 1), (2, 4, 4, 2), + (4, 4, 1, 1), (4, 4, 4, 4), (8, 8, 16, 4)): + with self.subTest(parallelism=parallelism, capacity=capacity, replicas=replicas): + collator = self.collator(capacity) + created = [] + barrier = threading.Barrier(expected) + + class Decoder: + def __init__(self, stream): + self.stream = stream + self.closed = False + created.append(self) + + def get_frames_at(self, indices): + barrier.wait(timeout=10) + # Non-contiguous decoder output exercises layout conversion. + data = torch.tensor(indices, dtype=torch.uint8)[:, None, None, None] + return SimpleNamespace(data=data.expand(-1, 4, 5, 3).permute(0, 3, 1, 2)) + + def close(self): + self.closed = True + + collator.decoder_factory = Decoder + try: + with patch("pypaimon.multimodal.video.Blob.from_file", self.streams), \ + patch("torch.get_num_threads", return_value=1): + for return_batch in (False, True): + for uint8 in (True, False): + result = _decode_video_windows( + plans, rows, [collator], features, uint8, + return_batch=return_batch, share_memory=return_batch, + decode_parallelism=parallelism, + decoders_per_payload=replicas)["camera"] + target = torch.tensor(windows, dtype=torch.uint8)[:, :, None, None, None] + target = target.expand(-1, -1, 3, 4, 5) + if not uint8: + target = target.float() / 255 + actual = result if return_batch else torch.stack(result) + self.assertTrue(torch.equal(actual, target)) + self.assertTrue(actual.is_contiguous()) + if return_batch: + self.assertTrue(result.is_shared()) + result[0].zero_() + self.assertTrue(torch.equal(result[2], target[2])) + self.assertEqual(len(created), expected) + self.assertEqual(len({id(d.stream) for d in created}), expected) + finally: + collator.close() + self.assertTrue(all(d.closed and d.stream.closed for d in created)) + + def test_single_payload_split_bounds_duplicate_frame_requests(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + from pypaimon.multimodal.lerobot.dataset import _decode_video_windows + from pypaimon.table.row.blob import VideoFrameDescriptor + + features = {"camera": {"dtype": "video", "shape": [4, 5, 3]}} + + def decode(windows): + rows = {i: {"camera": VideoFrameDescriptor( + "one-video", 0, 1, i, -1, 0).serialize()} + for i in {index for window in windows for index in window}} + plans = [{"windows": {"camera": window}} for window in windows] + collator = self.collator(4) + calls = [] + + class Decoder: + def __init__(self, stream): + self.stream = stream + + def get_frames_at(self, indices): + calls.append(tuple(indices)) + data = torch.tensor( + indices, dtype=torch.uint8)[:, None, None, None] + return SimpleNamespace( + data=data.expand(-1, 4, 5, 3).permute(0, 3, 1, 2)) + + def close(self): + pass + + collator.decoder_factory = Decoder + try: + with patch("pypaimon.multimodal.video.Blob.from_file", self.streams), \ + patch("torch.get_num_threads", return_value=1): + result = _decode_video_windows( + plans, rows, [collator], features, True, + return_batch=True, decode_parallelism=4, + decoders_per_payload=4)["camera"] + target = torch.tensor( + windows, dtype=torch.uint8)[:, :, None, None, None] + target = target.expand(-1, -1, 3, 4, 5) + self.assertTrue(torch.equal(result, target)) + return calls + finally: + collator.close() + + repeated = decode([[0, 1, 2]] * 4) + self.assertEqual(repeated, [(0, 1, 2)]) + + overlapping = decode([list(range(i, i + 8)) for i in range(8)]) + self.assertEqual(len(overlapping), 2) + self.assertEqual(sum(len(indices) for indices in overlapping), 22) + self.assertEqual(set().union(*map(set, overlapping)), set(range(15))) + + def test_replica_cache_eviction_and_pickle(self): + collator = self.collator(2) + descriptor = BlobDescriptor("one-video", 0, 1) + with patch("pypaimon.multimodal.video.Blob.from_file", self.streams): + first = collator._decoder(descriptor, None, 0) + second = collator._decoder(descriptor, None, 1) + self.assertIs(collator._decoder(descriptor), first) + third = collator._decoder(descriptor, None, 2) + self.assertTrue(second.closed) + self.assertFalse(first.closed) + child = pickle.loads(pickle.dumps(collator)) + self.assertFalse(child._decoders) + child.close() + collator.close() + self.assertTrue(first.closed and third.closed) def test_options(self): + self.assertEqual(CoreOptions(Options({})).read_video_decoders_per_payload(), 4) + self.assertEqual(CoreOptions(Options({"read.video.decoders-per-payload": "1"})) + .read_video_decoders_per_payload(), 1) + for value in (0, -1, True, 1.5, "no", "1.5"): + with self.subTest(replicas=value), self.assertRaises(ValueError): + CoreOptions(Options({"read.video.decoders-per-payload": value})) \ + .read_video_decoders_per_payload() self.assertEqual(CoreOptions(Options({})).read_video_decode_parallelism(), 16) self.assertEqual(CoreOptions(Options({"read.video.decode-parallelism": "1"})) .read_video_decode_parallelism(), 1)