This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 03d68bfec72e004731fbbc2956eadbadd566c8db Author: XiaoHongbo <[email protected]> AuthorDate: Wed Sep 16 15:21:41 2026 +0800 [python] Speed up LeRobot video window reads (#9852) --- .../pypaimon/multimodal/lerobot/dataset.py | 118 +++++++++++- .../pypaimon/tests/multimodal_lerobot_test.py | 203 +++++++++++++++++++++ 2 files changed, 315 insertions(+), 6 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 1332b870cf..eca82cbfd6 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -65,7 +65,7 @@ _TORCH_DTYPE_NAMES = { } _IMAGE_READ_ATTEMPTS = 3 -_MAX_VIDEO_DECODE_WORKERS = 8 +_MAX_VISUAL_WORKERS = 8 _CONTROL_FEATURES = frozenset({ "index", @@ -424,6 +424,14 @@ class PaimonDatasetReader(ABC): self.return_uint8, ) + # Transforms may shrink frames; assemble those windows one sample at a time. + video_windows = {} if self.image_transforms is not None else _decode_video_windows( + plans, rows, getattr(self, "_video_collators", ()), + self._features, self.return_uint8) + for group in row_groups: + for row in group.values(): + for key in video_windows: + row.pop(key, None) _decode_video_rows( row_groups, getattr(self, "_video_collators", ())) converted = { @@ -438,9 +446,14 @@ class PaimonDatasetReader(ABC): }) import torch + visual_windows = _stack_visual_windows( + plans, converted, [key for key in self._visual_keys + if key not in video_windows] + ) if self.image_transforms is None and plans[0]["windows"] else {} + visual_windows.update(video_windows) duplicates = _duplicate_indices(plans) result = [] - for plan in plans: + for offset, plan in enumerate(plans): item = dict(converted[plan["index"]]) if plan["index"] in duplicates: item = { @@ -448,9 +461,12 @@ class PaimonDatasetReader(ABC): for key, value in item.items() } for key, positions in plan["windows"].items(): - item[key] = torch.stack([ - converted[position][key] for position in positions - ]) + if key in visual_windows: + item[key] = visual_windows[key][offset] + else: + item[key] = torch.stack([ + converted[position][key] for position in positions + ]) item.update(plan["padding"]) if self.image_transforms is not None: for key in self._visual_keys: @@ -1309,6 +1325,96 @@ def _attach_task_labels(rows, task_names, subtask_names): row["subtask"] = subtask_names[subtask_index] +def _decode_video_windows(plans, rows, collators, features, return_uint8): + import torch + + tasks = [c for c in collators if c.video_column in plans[0]["windows"]] + if not tasks or torch.get_num_threads() > 1: + return {} + grad_enabled = torch.is_grad_enabled() + inference_enabled = torch.is_inference_mode_enabled() + + def decode(collator): + collator._ensure_process_local_cache() + key = collator.video_column + descriptors = { + position: collator._prepare_row(row)[1] + for position, row in rows.items() if key in row + } + requests = {} + for offset, plan in enumerate(plans): + window = [descriptors[p] for p in plan["windows"][key]] + if not window or any(d is None for d in window): + return key, None + payload = window[0].payload_descriptor + # Cross-file windows use the regular frame assembly path. + if any(d.payload_descriptor != payload for d in window): + return key, None + requests.setdefault(payload, []).append(( + offset, [d.frame_index for d in window])) + + output = [None] * len(plans) + with torch.inference_mode(inference_enabled), \ + torch.set_grad_enabled(grad_enabled): + for payload, windows in requests.items(): + decoder = collator._decoder(payload) + get_frames = getattr(decoder, "get_frames_at", None) + if not callable(get_frames): + return key, None + indices = sorted({i for _, window in windows for i in window}) + frames = get_frames(indices=indices).data + if (not torch.is_tensor(frames) or frames.ndim != 4 + or len(frames) != len(indices)): + raise ValueError("Video decoder must return one frame per index.") + _video_tensor(frames[0], features[key], return_uint8=True) + # Reorder reused frames once instead of once per window. + if sum(len(window) for _, window in windows) >= 2 * len(indices): + frames = frames.contiguous() + positions = {index: pos for pos, index in enumerate(indices)} + for offset, window in windows: + selection = [positions[index] for index in window] + start = selection[0] + if selection == list(range(start, start + len(selection))): + output[offset] = frames[start:start + len(selection)].clone( + memory_format=torch.contiguous_format) + else: + output[offset] = frames.index_select(0, torch.tensor( + selection, dtype=torch.long, device=frames.device)) + if frames.dtype == torch.uint8 and not return_uint8: + output[offset] = output[offset].float().div_(255) + return key, output + + if len(tasks) == 1: + decoded = [decode(tasks[0])] + else: + with ThreadPoolExecutor( + max_workers=min(len(tasks), _MAX_VISUAL_WORKERS)) as executor: + decoded = list(executor.map(decode, tasks)) + return {key: windows for key, windows in decoded if windows is not None} + + +def _stack_visual_windows(plans, rows, visual_keys): + import torch + + keys = [key for key in plans[0]["windows"] if key in visual_keys] + # Let PyTorch handle parallelism when its own thread pool is enabled. + if len(keys) < 2 or torch.get_num_threads() > 1: + return {} + grad_enabled = torch.is_grad_enabled() + inference_enabled = torch.is_inference_mode_enabled() + + def stack(key): + with torch.inference_mode(inference_enabled), \ + torch.set_grad_enabled(grad_enabled): + return [torch.stack([ + rows[position][key] for position in plan["windows"][key] + ]) for plan in plans] + + with ThreadPoolExecutor( + max_workers=min(len(keys), _MAX_VISUAL_WORKERS)) as executor: + return dict(zip(keys, executor.map(stack, keys))) + + def _torch_row(row, features, return_uint8=False): import torch @@ -1568,7 +1674,7 @@ def _decode_video_rows(row_groups, collators): else: with ThreadPoolExecutor( max_workers=min( - len(tasks), _MAX_VIDEO_DECODE_WORKERS)) as executor: + len(tasks), _MAX_VISUAL_WORKERS)) as executor: decoded_groups = list(executor.map( lambda task: task[0](task[2]), tasks)) diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index a9d0ff6533..7c07c0718a 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -26,6 +26,7 @@ import sys import tempfile import threading import unittest +import weakref from concurrent.futures import ThreadPoolExecutor from pathlib import Path from types import SimpleNamespace @@ -48,10 +49,12 @@ from pypaimon.multimodal.lerobot.dataset import ( _PyAVVideoDecoder, _arrow_rows, _decode_video_rows, + _decode_video_windows, _image_tensor, _index_names, _open_video_decoder, _selected_episodes, + _stack_visual_windows, _torch_row, ) from pypaimon.multimodal.lerobot.api import _create_target_table @@ -949,6 +952,206 @@ class LeRobotValidationTest(unittest.TestCase): self.assertEqual([3, 4, 5], list( sample["observation.image"].shape)) + def test_transformed_windows_are_assembled_per_sample(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + + keys = ["left", "right"] + original_stack = torch.stack + for kind in ("image", "video"): + with self.subTest(kind=kind): + reader = object.__new__(_ManualDatasetReader) + reader._total_frames = 33 + reader.episodes = None + reader._selected_ranges = None + reader._features = { + key: {"dtype": kind, "shape": [4, 5, 3]} + for key in keys + } + reader._projection = keys + reader._delta_projection = keys + reader._image_keys = [] + reader._visual_keys = keys + reader._video_collators = [] + reader._file_io = None + reader._task_names = ["task"] + reader._subtask_names = None + reader.return_uint8 = True + reader._plan = lambda i: { + "index": i, "windows": {key: [i, i, i + 1] for key in keys}, + "padding": {}, + } + reader._read_rows = lambda indices, projection: { + i: dict(task_index=0, **{ + key: torch.full((3, 4, 5), i, dtype=torch.uint8) + for key in keys + }) for i in indices + } + windows = [] + peak = [] + + def stack(*args, **kwargs): + value = original_stack(*args, **kwargs) + windows.append(weakref.ref(value)) + peak.append(sum(ref() is not None for ref in windows)) + return value + + def resize(value): + return torch.nn.functional.interpolate( + value.float(), size=(2, 2), mode="nearest") + + reader.image_transforms = resize + module = "pypaimon.multimodal.lerobot.dataset." + with patch("torch.get_num_threads", return_value=1), \ + patch("torch.stack", side_effect=stack), \ + patch(module + "_decode_video_windows", + wraps=_decode_video_windows) as decode, \ + patch(module + "_stack_visual_windows", + wraps=_stack_visual_windows) as assemble: + result = reader.get_items(list(range(32))) + self.assertEqual(64, len(windows)) + self.assertLessEqual(max(peak), len(keys)) + decode.assert_not_called() + assemble.assert_not_called() + self.assertTrue(all(ref() is None for ref in windows)) + for i, item in enumerate(result): + for key in keys: + expected = torch.tensor([i, i, i + 1]).reshape(3, 1, 1, 1) + self.assertTrue(torch.equal( + item[key], expected.expand(3, 3, 2, 2).float())) + result[0][keys[0]].zero_() + self.assertEqual(1, result[0][keys[1]][-1, 0, 0, 0].item()) + + def test_video_windows_decode_directly_and_fall_back(self): + try: + import torch + except ImportError as error: + self.skipTest(str(error)) + from pypaimon.multimodal.video import VideoFrameCollator + from pypaimon.table.row.blob import VideoFrameDescriptor + + key = "camera" + feature = {key: {"dtype": "video", "shape": [4, 5, 3]}} + rows = {i: {key: VideoFrameDescriptor("a.video", 0, 1, i).serialize()} + for i in range(2)} + plans = [{"windows": {key: [1, 0, 0]}}] * 2 + pixels = torch.arange(120).reshape(2, 4, 5, 3).to(torch.uint8) + frames = pixels.permute(0, 3, 1, 2) + batch = frames.contiguous() + decoder = SimpleNamespace(get_frames_at=Mock( + return_value=SimpleNamespace(data=batch))) + collator = VideoFrameCollator( + SimpleNamespace(file_io=Mock()), video_column=key, + decoder_factory=Mock(), decode_fn=Mock()) + with patch.object(collator, "_decoder", return_value=decoder) as open_decoder, \ + patch("torch.get_num_threads", return_value=1): + for uint8, batch in ((True, frames), (False, frames), + (True, frames.contiguous()), + (False, frames.contiguous())): + decoder.get_frames_at.reset_mock() + decoder.get_frames_at.return_value = SimpleNamespace(data=batch) + result = _decode_video_windows( + plans, rows, [collator], feature, uint8)[key] + expected = batch[[1, 0, 0]] + if not uint8: + expected = expected.float().div(255) + self.assertTrue(torch.equal(result[0], expected)) + self.assertTrue(result[0].is_contiguous()) + result[0][1].zero_() + self.assertTrue(torch.equal(result[0][2], expected[2])) + self.assertTrue(torch.equal(result[1], expected)) + self.assertTrue(torch.equal(batch, frames)) + decoder.get_frames_at.assert_called_once_with(indices=[0, 1]) + + for indices in ([0, 1], [1, 0, 0]): + decoder.get_frames_at.reset_mock() + decoder.get_frames_at.return_value = SimpleNamespace(data=frames) + result = _decode_video_windows( + [{"windows": {key: indices}}], rows, [collator], feature, True)[key][0] + decoder.get_frames_at.assert_called_once_with(indices=[0, 1]) + self.assertTrue(torch.equal(result, batch[indices])) + self.assertTrue(result.is_contiguous()) + result.zero_() + self.assertTrue(torch.equal(batch, frames)) + + decoder.get_frames_at.reset_mock() + separate = dict(rows) + separate[2] = {key: VideoFrameDescriptor("b.video", 0, 1, 0).serialize()} + separate[3] = {key: VideoFrameDescriptor("b.video", 0, 1, 1).serialize()} + interleaved = [{"windows": {key: window}} + for window in ([1, 0, 0], [2, 3], [0, 1], [3, 2, 2])] + result = _decode_video_windows( + interleaved, separate, [collator], feature, True)[key] + self.assertEqual(2, decoder.get_frames_at.call_count) + self.assertEqual( + [[0, 1], [0, 1]], + [c.kwargs["indices"] for c in decoder.get_frames_at.call_args_list]) + for actual, indices in zip(result, ([1, 0, 0], [0, 1], [0, 1], [1, 0, 0])): + self.assertTrue(torch.equal(actual, batch[indices])) + self.assertTrue(actual.is_contiguous()) + result[2].zero_() + self.assertTrue(torch.equal(result[1], batch)) + self.assertTrue(torch.equal(batch, frames)) + + open_decoder.reset_mock() + mixed = dict(rows) + mixed[1] = {key: VideoFrameDescriptor("b.video", 0, 1, 0).serialize()} + self.assertEqual({}, _decode_video_windows( + plans, mixed, [collator], feature, True)) + open_decoder.assert_not_called() + self.assertEqual({}, _decode_video_windows( + plans, {0: {key: None}, 1: rows[1]}, [collator], feature, True)) + open_decoder.return_value = SimpleNamespace() + self.assertEqual({}, _decode_video_windows( + plans, rows, [collator], feature, True)) + open_decoder.return_value = decoder + decoder.get_frames_at.return_value = SimpleNamespace(data=batch[:1]) + with self.assertRaisesRegex(ValueError, "one frame per index"): + _decode_video_windows(plans, rows, [collator], feature, True) + + def test_visual_windows_preserve_order_padding_and_isolation(self): + try: + import torch + from torch.utils.data import default_collate + except ImportError as error: + self.skipTest(str(error)) + + keys = ["left", "right"] + plans = [{"windows": {key: indices for key in keys}} + for indices in ([1, 0, 0], [0, 1, 1], [1, 0, 0])] + for dtype in (torch.uint8, torch.float32): + frames = torch.arange(120).reshape(2, 4, 5, 3).to(dtype) + frames = frames.permute(0, 3, 1, 2) + rows = {i: {key: frames[i] for key in keys} for i in range(2)} + with patch("torch.get_num_threads", return_value=1): + actual = _stack_visual_windows(plans, rows, keys) + for key in keys: + expected = [torch.stack([rows[i][key] for i in + plan["windows"][key]]) for plan in plans] + self.assertTrue(torch.equal( + default_collate(actual[key]), default_collate(expected))) + self.assertTrue(all(value.is_contiguous() + for value in actual[key])) + actual[key][0][1].zero_() + self.assertTrue(torch.equal(actual[key][0][2], expected[0][2])) + self.assertTrue(torch.equal(actual[key][2], expected[2])) + self.assertTrue(torch.equal(frames[0], expected[0][1])) + with patch("torch.get_num_threads", return_value=2): + self.assertEqual({}, _stack_visual_windows(plans, rows, keys)) + self.assertEqual({}, _stack_visual_windows(plans, rows, ["left"])) + + rows = {i: {key: torch.ones(2, requires_grad=True) for key in keys} + for i in range(2)} + with patch("torch.get_num_threads", return_value=1), torch.no_grad(): + result = _stack_visual_windows(plans, rows, keys) + self.assertFalse(result["left"][0].requires_grad) + with patch("torch.get_num_threads", return_value=1), \ + torch.inference_mode(): + result = _stack_visual_windows(plans, rows, keys) + self.assertTrue(result["left"][0].is_inference()) + def test_dataset_return_uint8_requires_bool(self): loaded = ( Mock(),
