This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new aafc7200a4 [python] Speed up LeRobot video window reads (#9852)
aafc7200a4 is described below
commit aafc7200a4dd4db8a4259487aa83d4e1d80cf19d
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(),
