From 853441a65fc67eb7b173e6bb07fe86de89a5d740 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 21 Oct 2026 03:20:31 -0700 Subject: [PATCH 1/2] [python] Allocate shared video batches without copying empty storage --- .../pypaimon/multimodal/lerobot/dataset.py | 18 +++++++-- .../tests/video_decode_parallelism_test.py | 38 +++++++++++++++++++ 2 files changed, 52 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 1ce3998ae47a..4ec0f3b3349f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1580,10 +1580,20 @@ def read_frames(decoder, windows): 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 = torch.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 17fa76a12446..7b57ffe1e48f 100644 --- a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py +++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py @@ -391,6 +391,44 @@ def close(self): 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))) + 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 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.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: + collator.close() + def test_real_pyav_payloads_match_serial(self): try: import av From e00f272f698dea7f44788ed8fb4cb6c7010ea856 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Wed, 21 Oct 2026 09:24:13 -0700 Subject: [PATCH 2/2] [python] Preserve decoder device when allocating shared batches --- .../pypaimon/multimodal/lerobot/dataset.py | 2 +- .../tests/video_decode_parallelism_test.py | 33 +++++++++++-------- 2 files changed, 20 insertions(+), 15 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 4ec0f3b3349f..e967b255c4fd 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1581,7 +1581,7 @@ def read_frames(decoder, windows): + tuple(frames.shape[1:])) if output is None: if share_memory and frames.device.type == "cpu": - empty = torch.empty(0, dtype=output_dtype) + empty = frames.new_empty(0, dtype=output_dtype) numel = 1 for size in shape: numel *= size diff --git a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py index 7b57ffe1e48f..dfeb41bbb57c 100644 --- a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py +++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py @@ -407,26 +407,31 @@ def test_single_payload_shared_batch_matches_windows(self): 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 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.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)) + 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):