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 729b62e2fb [python] Avoid an extra shared-memory copy for LeRobot
video batches (#10733)
729b62e2fb is described below
commit 729b62e2fbd8a2514ce376d1183c1e1830c50adf
Author: XiaoHongbo <[email protected]>
AuthorDate: Thu Oct 22 09:06:59 2026 +0800
[python] Avoid an extra shared-memory copy for LeRobot video batches
(#10733)
---
.../pypaimon/multimodal/lerobot/dataset.py | 18 +++++++--
.../tests/video_decode_parallelism_test.py | 43 ++++++++++++++++++++++
2 files changed, 57 insertions(+), 4 deletions(-)
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 5456ff7cf7..60386c82f4 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -1580,10 +1580,20 @@ def _decode_video_windows_with_pool(
shape = ((len(plans), window_size)
+ tuple(frames.shape[1:]))
if output is None:
- output = torch.empty(
- shape, dtype=output_dtype, device=frames.device)
- if share_memory:
- output.share_memory_()
+ if share_memory and frames.device.type == "cpu":
+ empty = frames.new_empty(0, dtype=output_dtype)
+ numel = 1
+ for size in shape:
+ numel *= size
+ # Match DataLoader collation: allocate shared
storage directly.
+ storage = empty._typed_storage()._new_shared(
+ numel, device=empty.device)
+ output = empty.new(storage).resize_(shape)
+ else:
+ output = torch.empty(
+ shape, dtype=output_dtype,
device=frames.device)
+ if share_memory:
+ output.share_memory_()
elif tuple(output.shape) != shape \
or output.dtype != output_dtype:
raise ValueError(
diff --git a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
index 17fa76a124..dfeb41bbb5 100644
--- a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
+++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
@@ -391,6 +391,49 @@ class VideoDecodeParallelismTest(unittest.TestCase):
finally:
dataset.close()
+ def test_single_payload_shared_batch_matches_windows(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
+
+ pixels = torch.arange(10 * 4 * 5 * 3).reshape(10, 4, 5,
3).to(torch.uint8)
+ rows = {i: {"camera": VideoFrameDescriptor("video", 0, 1, i, -1,
0).serialize()}
+ for i in range(10)}
+ windows = ([0, 0, 2], [2, 3, 4], [9, 7, 8], [0, 0, 2])
+ plans = [{"windows": {"camera": indices}} for indices in windows]
+ collator = self.collator()
+ decoder = SimpleNamespace(get_frames_at=lambda indices:
SimpleNamespace(
+ data=pixels[indices].permute(0, 3, 1, 2)))
+ original_device = torch.get_default_device()
+ try:
+ with patch.object(collator, "_decoder", return_value=decoder), \
+ patch("torch.get_num_threads", return_value=1), \
+ patch.object(torch.Tensor, "share_memory_",
+ side_effect=AssertionError("Unexpected
storage copy")):
+ for default_device in ("cpu", "meta"):
+ torch.set_default_device(default_device)
+ for uint8 in (True, False):
+ result = _decode_video_windows(
+ plans, rows, [collator],
+ {"camera": {"dtype": "video", "shape": [4, 5,
3]}}, uint8,
+ return_batch=True, share_memory=True)["camera"]
+ expected = torch.stack([pixels[w].permute(0, 3, 1, 2)
for w in windows])
+ if not uint8:
+ expected = expected.float().div_(255)
+ self.assertEqual(result.device.type, "cpu")
+ self.assertTrue(torch.equal(result, expected))
+ self.assertTrue(result.is_shared())
+ self.assertTrue(result.is_contiguous())
+ saved = result[3].clone()
+ result[0].zero_()
+ self.assertTrue(torch.equal(result[3], saved))
+ finally:
+ torch.set_default_device(original_device)
+ collator.close()
+
def test_real_pyav_payloads_match_serial(self):
try:
import av
