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 2620ab68c1 [Python] Decode LeRobot video payloads in parallel (#10703)
2620ab68c1 is described below
commit 2620ab68c148a869e0aa91d899e41b72fc99f2cd
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Oct 21 20:03:21 2026 +0800
[Python] Decode LeRobot video payloads in parallel (#10703)
---
docs/docs/pypaimon/lerobot.md | 3 +
.../pypaimon/common/options/core_options.py | 16 +-
.../pypaimon/multimodal/lerobot/dataset.py | 96 +++--
paimon-python/pypaimon/multimodal/video.py | 72 +++-
.../pypaimon/tests/multimodal_lerobot_test.py | 2 +-
.../tests/video_decode_parallelism_test.py | 455 +++++++++++++++++++++
6 files changed, 618 insertions(+), 26 deletions(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index 45120edb2f..ae7ca1910a 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -256,6 +256,9 @@ Reads stay on the selected snapshot; unavailable or
unsupported native paths use
Set `table = table.copy({"read.video.max-open-decoders": "8"})` before
constructing the dataset to limit cached decoders per camera per worker
(default: 16).
Each open decoder uses memory and a file handle, so usage grows with cameras
and workers.
+`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.
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 a9c544303e..77f40584ee 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -1172,6 +1172,15 @@ class CoreOptions:
"LeRobot DataLoader worker.")
)
+ READ_VIDEO_DECODE_PARALLELISM: ConfigOption[int] = (
+ ConfigOptions.key("read.video.decode-parallelism")
+ .int_type()
+ .default_value(16)
+ .with_description(
+ "Maximum concurrent video payload decodes shared by all cameras in
a "
+ "PyPaimon LeRobot Dataset worker.")
+ )
+
PARQUET_COLUMN_INDEX_ENABLED: ConfigOption[bool] = (
ConfigOptions.key("parquet.filter.columnindex.enabled")
.boolean_type()
@@ -1937,7 +1946,12 @@ class CoreOptions:
return self.options.get(CoreOptions.READ_BATCH_SIZE, default or 1024)
def read_video_max_open_decoders(self) -> int:
- option = CoreOptions.READ_VIDEO_MAX_OPEN_DECODERS
+ return
self._positive_video_option(CoreOptions.READ_VIDEO_MAX_OPEN_DECODERS)
+
+ def read_video_decode_parallelism(self) -> int:
+ return
self._positive_video_option(CoreOptions.READ_VIDEO_DECODE_PARALLELISM)
+
+ def _positive_video_option(self, option) -> int:
raw = self.options.to_map().get(option.key())
if isinstance(raw, (bool, float)):
raise ValueError("%s must be a positive integer." % option.key())
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 22a0fb65c9..21d4bb8b48 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -50,7 +50,7 @@ from pypaimon.multimodal.lerobot.schema import (
_validate_lerobot_schema,
)
from pypaimon.multimodal.table import _target_schema, _time_travel_table
-from pypaimon.multimodal.video import VideoFrameCollator
+from pypaimon.multimodal.video import VideoFrameCollator, _VideoDecodePool
from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex
@@ -274,6 +274,12 @@ class PaimonDatasetReader(ABC):
if isinstance(table_options, CoreOptions)
else CoreOptions.READ_VIDEO_MAX_OPEN_DECODERS.default_value()
)
+ self._video_decode_parallelism = (
+ table_options.read_video_decode_parallelism()
+ if isinstance(table_options, CoreOptions)
+ else CoreOptions.READ_VIDEO_DECODE_PARALLELISM.default_value()
+ )
+ self._video_decode_pool = _VideoDecodePool()
self._video_collators = [
VideoFrameCollator(
access,
@@ -448,13 +454,21 @@ class PaimonDatasetReader(ABC):
video_windows = {} if self.image_transforms is not None else
_decode_video_windows(
plans, rows, getattr(self, "_video_collators", ()),
self._features, self.return_uint8,
- return_batch=return_batch, share_memory=share_memory)
+ return_batch=return_batch, share_memory=share_memory,
+ decode_pool=getattr(self, "_video_decode_pool", None),
+ decode_parallelism=getattr(
+ self, "_video_decode_parallelism",
+ CoreOptions.READ_VIDEO_DECODE_PARALLELISM.default_value()))
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", ()))
+ row_groups, getattr(self, "_video_collators", ()),
+ decode_pool=getattr(self, "_video_decode_pool", None),
+ decode_parallelism=getattr(
+ self, "_video_decode_parallelism",
+ CoreOptions.READ_VIDEO_DECODE_PARALLELISM.default_value()))
converted = {
position: _torch_row(
row, self._features, self.return_uint8)
@@ -509,6 +523,9 @@ class PaimonDatasetReader(ABC):
return self.get_items(indices)
def close(self):
+ pool = getattr(self, "_video_decode_pool", None)
+ if pool is not None:
+ pool.close()
first_error = None
for collator in getattr(self, "_video_collators", ()):
try:
@@ -1422,7 +1439,21 @@ 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):
+ return_batch=False, share_memory=False, decode_parallelism=1,
decode_pool=None):
+ 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))
+ finally:
+ if decode_pool is None:
+ pool.close()
+
+
+def _decode_video_windows_with_pool(
+ plans, rows, collators, features, return_uint8, *,
+ return_batch, share_memory, decode_parallelism, executor):
import torch
tasks = [c for c in collators if c.video_column in plans[0]["windows"]]
@@ -1456,16 +1487,24 @@ def _decode_video_windows(
requests.setdefault((payload, keyframe_index), []).append((
offset, [d.frame_index for d in window]))
+ def read_frames(decoder, windows):
+ get_frames = getattr(decoder, "get_frames_at", None)
+ if not callable(get_frames):
+ return None
+ indices = sorted({i for _, window in windows for i in window})
+ with torch.inference_mode(inference_enabled), \
+ torch.set_grad_enabled(grad_enabled):
+ return windows, indices, get_frames(indices=indices).data
+
output = None if return_batch else [None] * len(plans)
with torch.inference_mode(inference_enabled), \
- torch.set_grad_enabled(grad_enabled):
- for (payload, keyframe_index), windows in requests.items():
- decoder = collator._decoder(payload, keyframe_index)
- get_frames = getattr(decoder, "get_frames_at", None)
- if not callable(get_frames):
+ torch.set_grad_enabled(grad_enabled), \
+ collator._map_decoders(
+ requests, read_frames, decode_parallelism, executor) as
decoded:
+ for result in decoded:
+ if result is None:
return key, None
- indices = sorted({i for _, window in windows for i in window})
- frames = get_frames(indices=indices).data
+ windows, indices, frames = result
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.")
@@ -1519,12 +1558,12 @@ def _decode_video_windows(
output[offset] = output[offset].float().div_(255)
return key, output
- if len(tasks) == 1:
- decoded = [decode(tasks[0])]
+ if executor is None or len(tasks) == 1:
+ decoded = [decode(task) for task in tasks]
else:
with ThreadPoolExecutor(
- max_workers=min(len(tasks), _MAX_VISUAL_WORKERS)) as executor:
- decoded = list(executor.map(decode, tasks))
+ max_workers=min(len(tasks), _MAX_VISUAL_WORKERS)) as
camera_executor:
+ decoded = list(camera_executor.map(decode, tasks))
return {key: windows for key, windows in decoded if windows is not None}
@@ -2151,7 +2190,8 @@ def _identity(values):
return values
-def _decode_video_rows(row_groups, collators):
+def _decode_video_rows(
+ row_groups, collators, decode_parallelism=_MAX_VISUAL_WORKERS,
decode_pool=None):
tasks = []
for collator in collators:
targets = []
@@ -2167,14 +2207,24 @@ def _decode_video_rows(row_groups, collators):
if not tasks:
return
- if len(tasks) == 1:
- decoded_groups = [tasks[0][0](tasks[0][2])]
+ if decode_parallelism == 1:
+ decoded_groups = [collator(rows) for collator, _, rows in tasks]
else:
- with ThreadPoolExecutor(
- max_workers=min(
- len(tasks), _MAX_VISUAL_WORKERS)) as executor:
- decoded_groups = list(executor.map(
- lambda task: task[0](task[2]), tasks))
+ from concurrent.futures import wait
+
+ pool = decode_pool if decode_pool is not None else _VideoDecodePool()
+ futures = []
+ try:
+ executor = pool.get(decode_parallelism)
+ for collator, _, rows in tasks:
+ futures.append(executor.submit(collator, rows))
+ decoded_groups = [future.result() for future in futures]
+ finally:
+ for future in futures:
+ future.cancel()
+ wait(futures)
+ if decode_pool is None:
+ pool.close()
for (collator, targets, _), decoded in zip(tasks, decoded_groups):
for (rows, index), row in zip(targets, decoded):
diff --git a/paimon-python/pypaimon/multimodal/video.py
b/paimon-python/pypaimon/multimodal/video.py
index 90489a0b0a..3d9eb812c4 100644
--- a/paimon-python/pypaimon/multimodal/video.py
+++ b/paimon-python/pypaimon/multimodal/video.py
@@ -19,8 +19,10 @@
import io
import os
-from collections import OrderedDict
+from collections import OrderedDict, deque
from collections.abc import Mapping
+from concurrent.futures import ThreadPoolExecutor, wait
+from contextlib import contextmanager
from fractions import Fraction
from itertools import chain
@@ -28,6 +30,35 @@ from pypaimon.table.row.blob import Blob,
VideoFrameDescriptor
from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex
+class _VideoDecodePool:
+ """One lazily created decode pool per dataset and worker process."""
+
+ def __init__(self):
+ self._executor = None
+ self._parallelism = 0
+ self._pid = os.getpid()
+
+ def get(self, parallelism):
+ if self._pid != os.getpid():
+ # Never join threads inherited from the parent process.
+ self._executor = None
+ self._pid = os.getpid()
+ if self._executor is not None and self._parallelism != parallelism:
+ self.close()
+ if parallelism > 1 and self._executor is None:
+ self._executor = ThreadPoolExecutor(max_workers=parallelism)
+ self._parallelism = parallelism
+ return self._executor
+
+ def close(self):
+ executor, self._executor = self._executor, None
+ if executor is not None and self._pid == os.getpid():
+ executor.shutdown(wait=True)
+
+ def __getstate__(self):
+ return {"_executor": None, "_parallelism": 0, "_pid": None}
+
+
def _encode_images_to_video(
images,
output_path,
@@ -225,6 +256,45 @@ class VideoFrameCollator:
state["_owner_pid"] = None
return state
+ @contextmanager
+ def _map_decoders(self, requests, decode, parallelism, executor=None):
+ self._ensure_process_local_cache()
+ limit = min(parallelism, self.max_open_videos, len(requests))
+ if executor is None:
+ yield (decode(self._decoder(*key), value)
+ for key, value in requests.items())
+ return
+ items = iter(requests.items())
+ pending = deque()
+
+ def submit():
+ try:
+ key, value = next(items)
+ except StopIteration:
+ return
+ # Keep cache mutations ordered; pending decoders remain most
recent.
+ decoder = self._decoder(*key)
+ pending.append(executor.submit(decode, decoder, value))
+
+ def results():
+ while pending:
+ future = pending[0]
+ value = future.result()
+ pending.popleft()
+ yield value
+ del value, future
+ submit()
+
+ try:
+ for _ in range(limit):
+ submit()
+ yield results()
+ finally:
+ for future in pending:
+ future.cancel()
+ # A failed or abandoned batch must finish before cache reuse/close.
+ wait(pending)
+
def _decode_rows(self, rows):
decoded = [None] * len(rows)
grouped = OrderedDict()
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index e3cfc4f6da..acd87cccb0 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -266,7 +266,7 @@ class LeRobotValidationTest(unittest.TestCase):
0: {"camera_%d" % index: index for index in range(9)},
}
with patch(
- "pypaimon.multimodal.lerobot.dataset.ThreadPoolExecutor",
+ "pypaimon.multimodal.video.ThreadPoolExecutor",
side_effect=executor):
_decode_video_rows(
[rows], [Collator(index) for index in range(9)])
diff --git a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
new file mode 100644
index 0000000000..45a27930d1
--- /dev/null
+++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
@@ -0,0 +1,455 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import io
+import multiprocessing
+import os
+import pickle
+import threading
+import unittest
+from collections import OrderedDict
+from types import SimpleNamespace
+from unittest.mock import patch
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.common.options.options import Options
+from pypaimon.multimodal.video import VideoFrameCollator, _VideoDecodePool
+from pypaimon.table.row.blob import BlobDescriptor
+
+
+def identity(value):
+ return value
+
+
+def decode_worker(collator, pool, connection):
+ try:
+ requests = OrderedDict(
+ ((BlobDescriptor("child-%s" % i, 0, 1), None), i)
+ for i in range(2))
+ with patch("pypaimon.multimodal.video.Blob.from_file",
side_effect=lambda *args:
+ SimpleNamespace(new_input_stream=lambda:
io.BytesIO(b"video"))):
+ with collator._map_decoders(
+ requests, lambda stream, value: value, 2, pool.get(2)) as
values:
+ result = list(values)
+ connection.send((result, collator._owner_pid == os.getpid()))
+ finally:
+ pool.close()
+ collator.close()
+ connection.close()
+
+
+class VideoDecodeParallelismTest(unittest.TestCase):
+ def collator(self, capacity=2):
+ return VideoFrameCollator(
+ SimpleNamespace(file_io="test-file-io"), video_column="camera",
+ decoder_factory=identity, decode_fn=identity,
+ max_open_videos=capacity)
+
+ def requests(self, count=8):
+ return OrderedDict(
+ ((BlobDescriptor("video-%s" % i, 0, 1), None), i)
+ for i in range(count))
+
+ def streams(self, *args):
+ return SimpleNamespace(new_input_stream=lambda: io.BytesIO(b"video"))
+
+ def test_shared_limit_across_cameras_and_eviction(self):
+ from concurrent.futures import ThreadPoolExecutor
+
+ pool = _VideoDecodePool()
+ executor = pool.get(2)
+ collators = [self.collator() for _ in range(3)]
+ active = [0, 0]
+ lock = threading.Lock()
+ barrier = threading.Barrier(2)
+
+ def decode(stream, value):
+ with lock:
+ active[0] += 1
+ active[1] = max(active)
+ try:
+ barrier.wait(timeout=5)
+ self.assertFalse(stream.closed)
+ return value
+ finally:
+ with lock:
+ active[0] -= 1
+
+ def camera(collator, count):
+ with collator._map_decoders(self.requests(count), decode, 2,
executor) as values:
+ return list(values)
+
+ try:
+ with patch("pypaimon.multimodal.video.Blob.from_file",
self.streams), \
+ ThreadPoolExecutor(max_workers=3) as cameras:
+ # One camera has a single payload; it must also share the
budget.
+ futures = [cameras.submit(camera, c, n)
+ for c, n in zip(collators, (1, 3, 8))]
+ self.assertEqual([f.result() for f in futures],
+ [list(range(n)) for n in (1, 3, 8)])
+ self.assertEqual(active, [0, 2])
+ self.assertIs(executor, pool.get(2))
+ finally:
+ pool.close()
+ for collator in collators:
+ streams = list(collator._decoders.values())
+ collator.close()
+ self.assertTrue(all(stream.closed for _, stream in streams))
+ self.assertTrue(executor._shutdown)
+
+ def test_error_drains_running_decodes(self):
+ pool = _VideoDecodePool()
+ executor = pool.get(2)
+ collator = self.collator()
+ barrier = threading.Barrier(2)
+ started = threading.Event()
+ release = threading.Event()
+ finished = threading.Event()
+ errors = []
+
+ def decode(stream, value):
+ barrier.wait(timeout=5)
+ if value == 0:
+ raise ValueError("decode failed")
+ started.set()
+ self.assertTrue(release.wait(5))
+ self.assertFalse(stream.closed)
+ return value
+
+ def consume():
+ try:
+ with collator._map_decoders(self.requests(2), decode, 2,
executor) as values:
+ list(values)
+ except Exception as error:
+ errors.append(error)
+ finally:
+ finished.set()
+
+ with patch("pypaimon.multimodal.video.Blob.from_file", self.streams):
+ thread = threading.Thread(target=consume)
+ thread.start()
+ try:
+ self.assertTrue(started.wait(5))
+ self.assertFalse(finished.is_set())
+ finally:
+ release.set()
+ thread.join(5)
+ pool.close()
+ collator.close()
+ self.assertTrue(finished.is_set())
+ self.assertEqual(str(errors[0]), "decode failed")
+
+ def test_serial_and_pickle(self):
+ pool = _VideoDecodePool()
+ self.assertIsNone(pool.get(1))
+ collator = self.collator(1)
+ owner = threading.get_ident()
+ with patch("pypaimon.multimodal.video.Blob.from_file", self.streams):
+ with collator._map_decoders(
+ self.requests(3), lambda stream, value:
threading.get_ident(), 1) as values:
+ self.assertEqual(list(values), [owner] * 3)
+ pool.get(2).submit(identity, 1).result()
+ restored_collator, restored_pool =
pickle.loads(pickle.dumps((collator, pool)))
+ self.assertIsNone(restored_pool._executor)
+ self.assertFalse(restored_collator._decoders)
+ pool.close()
+ collator.close()
+ restored_pool.close()
+ restored_collator.close()
+
+ @unittest.skipUnless(hasattr(os, "fork"), "requires fork")
+ def test_fork_discards_inherited_executor(self):
+ pool = _VideoDecodePool()
+ executor = pool.get(2)
+ try:
+ with patch("pypaimon.multimodal.video.os.getpid",
return_value=-1), \
+ patch.object(executor, "shutdown",
side_effect=AssertionError("inherited pool")):
+ child_executor = pool.get(2)
+ self.assertIsNot(executor, child_executor)
+ pool.close()
+ self.assertTrue(child_executor._shutdown)
+ finally:
+ executor.shutdown(wait=True)
+
+ def test_workers_rebuild_pool(self):
+ collator = self.collator()
+ pool = _VideoDecodePool()
+ pool.get(2).submit(identity, 1).result()
+ try:
+ for method in ("spawn", "fork"):
+ if method not in multiprocessing.get_all_start_methods():
+ continue
+ with self.subTest(method=method):
+ ctx = multiprocessing.get_context(method)
+ receive, send = ctx.Pipe(duplex=False)
+ worker = ctx.Process(target=decode_worker, args=(collator,
pool, send))
+ worker.start()
+ send.close()
+ try:
+ self.assertTrue(receive.poll(15), "worker did not
complete")
+ self.assertEqual(receive.recv(), ([0, 1], True))
+ worker.join(5)
+ self.assertEqual(worker.exitcode, 0)
+ finally:
+ if worker.is_alive():
+ worker.terminate()
+ worker.join(5)
+ receive.close()
+ finally:
+ pool.close()
+ collator.close()
+
+ def test_parallel_windows_match_serial_and_preserve_torch_mode(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(4 * 4 * 5 * 3).reshape(4, 4, 5,
3).to(torch.uint8)
+ frames = pixels.permute(0, 3, 1, 2)
+ rows = {i: {"camera": VideoFrameDescriptor(
+ "video-%s" % (i // 4), 0, 1, i % 4, -1, 0).serialize()}
+ for i in range(12)}
+ plans = [{"windows": {"camera": w}} for w in (
+ [0, 0, 2], [4, 5, 6], [8, 9, 10], [0, 1, 3])]
+ features = {"camera": {"dtype": "video", "shape": [4, 5, 3]}}
+ modes = []
+
+ def read(indices):
+ modes.append((torch.is_grad_enabled(),
torch.is_inference_mode_enabled()))
+ # Preserve the noncontiguous layout returned by TorchCodec.
+ return SimpleNamespace(data=pixels[indices].permute(0, 3, 1, 2))
+
+ collator = self.collator(2)
+ decoder = SimpleNamespace(get_frames_at=read)
+ try:
+ with patch.object(collator, "_decoder", return_value=decoder), \
+ patch("torch.get_num_threads", return_value=1), \
+ torch.inference_mode():
+ for uint8 in (True, False):
+ for batch in (True, False):
+ with self.subTest(uint8=uint8, batch=batch):
+ kwargs = dict(return_batch=batch,
share_memory=batch)
+ expected = _decode_video_windows(
+ plans, rows, [collator], features, uint8,
+ decode_parallelism=1, **kwargs)["camera"]
+ actual = _decode_video_windows(
+ plans, rows, [collator], features, uint8,
+ decode_parallelism=4, **kwargs)["camera"]
+ if not batch:
+ expected, actual = torch.stack(expected),
torch.stack(actual)
+ self.assertTrue(torch.equal(expected, actual))
+ self.assertTrue(actual.is_contiguous())
+ if batch:
+ self.assertTrue(actual.is_shared())
+ saved = actual[3].clone()
+ actual[0].zero_()
+ self.assertTrue(torch.equal(saved, actual[3]))
+ self.assertTrue(modes)
+ self.assertEqual(set(modes), {(False, True)})
+ self.assertTrue(torch.equal(frames, pixels.permute(0, 3, 1, 2)))
+ finally:
+ collator.close()
+
+ def test_video_windows_share_worker_pool_across_cameras(self):
+ try:
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+ from contextlib import ExitStack
+ from pypaimon.multimodal.lerobot.dataset import _decode_video_windows
+ from pypaimon.table.row.blob import VideoFrameDescriptor
+
+ pool = _VideoDecodePool()
+ collators = [self.collator(1) for _ in range(3)]
+ names = ["camera-%s" % i for i in range(3)]
+ for collator, name in zip(collators, names):
+ collator.video_column = name
+ rows = {0: {name: VideoFrameDescriptor(name, 0, 1, 0, -1,
0).serialize()
+ for name in names}}
+ plans = [{"windows": {name: [0] for name in names}}]
+ features = {name: {"dtype": "video", "shape": [4, 5, 3]} for name in
names}
+ threads = set()
+ lock = threading.Lock()
+
+ def read(indices):
+ with lock:
+ threads.add(threading.get_ident())
+ return SimpleNamespace(data=torch.ones((len(indices), 3, 4, 5),
dtype=torch.uint8))
+
+ try:
+ with ExitStack() as stack:
+ stack.enter_context(patch("torch.get_num_threads",
return_value=1))
+ for collator in collators:
+ stack.enter_context(patch.object(
+ collator, "_decoder",
return_value=SimpleNamespace(get_frames_at=read)))
+ for parallelism in (1, 2):
+ threads.clear()
+ result = _decode_video_windows(
+ plans, rows, collators, features, True,
return_batch=True,
+ decode_parallelism=parallelism, decode_pool=pool)
+ self.assertEqual(set(result), set(names))
+ self.assertLessEqual(len(threads), parallelism)
+ for value in result.values():
+ self.assertTrue(torch.equal(value,
torch.ones_like(value)))
+ if parallelism == 1:
+ self.assertEqual(threads, {threading.get_ident()})
+ else:
+ self.assertNotIn(threading.get_ident(), threads)
+ executor = pool.get(2)
+ _decode_video_windows(plans, rows, collators, features, True,
+ decode_parallelism=2, decode_pool=pool)
+ self.assertIs(executor, pool.get(2))
+ finally:
+ pool.close()
+ for collator in collators:
+ collator.close()
+
+ def test_dataloader_fallback_honors_decode_limit(self):
+ try:
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+ import time
+ import pyarrow as pa
+ import pypaimon.multimodal as pm
+
+ names = ("camera_a", "camera_b")
+ features = {name: {"dtype": "int64", "shape": [1]}
+ for name in ("index", "episode_index", "frame_index",
"task_index")}
+ features["timestamp"] = {"dtype": "float32", "shape": [1]}
+ features.update({name: {"dtype": "video", "shape": [2, 2, 3]} for name
in names})
+ metadata = {
+ "repo_id": "test/fallback",
+ "info": {"codebase_version": "v3.0", "fps": 10, "total_frames": 1,
+ "total_episodes": 1, "total_tasks": 1, "features":
features},
+ "episodes": [{"episode_index": 0, "dataset_from_index": 0,
+ "dataset_to_index": 1, "length": 1, "tasks":
["pick"]}],
+ "tasks": ["pick"], "stats": {},
+ }
+
+ class Reader(pm.PaimonDatasetReader):
+ def read_indices(self, indices, columns):
+ row = {name: 0 for name in features}
+ row.update({name: b"video" for name in names})
+ return pa.Table.from_pylist(
+ [{name: row[name] for name in columns}],
schema=self.schema)
+
+ active = [0, 0]
+ lock = threading.Lock()
+
+ class Collator:
+ def __init__(self, name):
+ self.video_column = self.output_column = name
+
+ def __call__(self, rows):
+ with lock:
+ active[0] += 1
+ active[1] = max(active)
+ try:
+ time.sleep(0.02)
+ return [dict(row, **{self.output_column: torch.ones(
+ (3, 2, 2), dtype=torch.uint8)}) for row in rows]
+ finally:
+ with lock:
+ active[0] -= 1
+
+ def close(self):
+ pass
+
+ for transforms in (None, identity):
+ with self.subTest(transforms=transforms):
+ reader = Reader(metadata, file_io="test-file-io",
+ delta_timestamps={name: [0.0] for name in
names},
+ image_transforms=transforms, return_uint8=True)
+ reader._video_collators = [Collator(name) for name in names]
+ reader._video_decode_parallelism = 1
+ dataset = pm.PaimonLeRobotDataset(reader)
+ active[:] = [0, 0]
+ try:
+ with patch("torch.get_num_threads", return_value=2):
+ batch = next(iter(dataset.to_dataloader(batch_size=1,
num_workers=0)))
+ self.assertEqual(active, [0, 1])
+ for name in names:
+ self.assertTrue(torch.equal(batch[name],
torch.ones_like(batch[name])))
+ finally:
+ dataset.close()
+
+ def test_real_pyav_payloads_match_serial(self):
+ try:
+ import av
+ import numpy as np
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+ from functools import partial
+ from pypaimon.multimodal.lerobot.dataset import (
+ _decode_video_windows, _open_video_decoder)
+ from pypaimon.table.row.blob import VideoFrameDescriptor
+
+ videos = {}
+ for video in range(3):
+ target = io.BytesIO()
+ with av.open(target, mode="w", format="mp4") as container:
+ stream = container.add_stream("mpeg4", rate=30)
+ stream.width, stream.height = 16, 16
+ stream.pix_fmt = "yuv420p"
+ for frame in range(8):
+ image = np.full((16, 16, 3), video * 60 + frame * 5,
dtype=np.uint8)
+ for packet in
stream.encode(av.VideoFrame.from_ndarray(image, format="rgb24")):
+ container.mux(packet)
+ for packet in stream.encode():
+ container.mux(packet)
+ videos["video-%s" % video] = target.getvalue()
+ rows = {i: {"camera": VideoFrameDescriptor(
+ "video-%s" % (i // 8), 0, len(videos["video-%s" % (i // 8)]),
+ i % 8, -1, 0).serialize()} for i in range(24)}
+ plans = [{"windows": {"camera": indices}} for indices in (
+ [0, 0, 3], [8, 9, 11], [16, 18, 19], [1, 2, 4])]
+ outputs = []
+
+ def blob(file_io, uri, offset, length):
+ return SimpleNamespace(new_input_stream=lambda:
io.BytesIO(videos[uri]))
+
+ for parallelism 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=2)
+ try:
+ with patch("pypaimon.multimodal.video.Blob.from_file", blob), \
+ patch("torch.get_num_threads", return_value=1):
+ outputs.append(_decode_video_windows(
+ plans, rows, [collator],
+ {"camera": {"dtype": "video", "shape": [16, 16, 3]}},
True,
+ return_batch=True,
decode_parallelism=parallelism)["camera"])
+ finally:
+ collator.close()
+ self.assertTrue(torch.equal(*outputs))
+
+ def test_options(self):
+
self.assertEqual(CoreOptions(Options({})).read_video_decode_parallelism(), 16)
+ self.assertEqual(CoreOptions(Options({"read.video.decode-parallelism":
"1"}))
+ .read_video_decode_parallelism(), 1)
+ self.assertEqual(CoreOptions(Options({"read.video.decode-parallelism":
"4"}))
+ .read_video_decode_parallelism(), 4)
+ for value in (0, -1, True, 1.5, "no", "1.5"):
+ with self.subTest(value=value), self.assertRaises(ValueError):
+ CoreOptions(Options({"read.video.decode-parallelism": value}))
\
+ .read_video_decode_parallelism()
