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 dab9a1801b [python] Make LeRobot PyAV codec threads configurable
(#10796)
dab9a1801b is described below
commit dab9a1801b830750539154ea48dedc461891d2ae
Author: XiaoHongbo <[email protected]>
AuthorDate: Sun Oct 25 10:45:53 2026 +0800
[python] Make LeRobot PyAV codec threads configurable (#10796)
---
docs/docs/pypaimon/lerobot.md | 4 ++
.../pypaimon/common/options/core_options.py | 21 ++++++++
.../pypaimon/multimodal/lerobot/dataset.py | 22 +++++---
.../pypaimon/tests/multimodal_lerobot_test.py | 6 ++-
.../tests/video_decode_parallelism_test.py | 62 ++++++++++++++++++++++
5 files changed, 108 insertions(+), 7 deletions(-)
diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md
index a3c0cfcc78..caca02fee7 100644
--- a/docs/docs/pypaimon/lerobot.md
+++ b/docs/docs/pypaimon/lerobot.md
@@ -261,6 +261,10 @@ Each open decoder uses memory and a file handle, so usage
grows with cameras and
`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.
+`read.video.codec-threads` controls FFmpeg threads per PyAV decoder (default:
1),
+avoiding nested codec parallelism with multiple workers and concurrent videos.
+Set it to 0 for automatic selection. It does not affect TorchCodec.
+
`read.video.decoders-per-payload` (default: 4) also splits single-video
windows across
independent cached decoders, within those limits. Set it to 1 to disable this
splitting;
the reader automatically reduces the decoder count when overlapping windows
would add
diff --git a/paimon-python/pypaimon/common/options/core_options.py
b/paimon-python/pypaimon/common/options/core_options.py
index 32b8fdb45a..d462943b93 100644
--- a/paimon-python/pypaimon/common/options/core_options.py
+++ b/paimon-python/pypaimon/common/options/core_options.py
@@ -1180,6 +1180,15 @@ class CoreOptions:
"PyPaimon LeRobot Dataset worker.")
)
+ READ_VIDEO_CODEC_THREADS: ConfigOption[int] = (
+ ConfigOptions.key("read.video.codec-threads")
+ .int_type()
+ .default_value(1)
+ .with_description(
+ "FFmpeg threads per PyAV video decoder in a LeRobot worker. "
+ "0 keeps automatic selection; 1 avoids nested codec parallelism.")
+ )
+
READ_VIDEO_DECODERS_PER_PAYLOAD: ConfigOption[int] = (
ConfigOptions.key("read.video.decoders-per-payload")
.int_type()
@@ -1962,6 +1971,18 @@ class CoreOptions:
def read_video_decoders_per_payload(self) -> int:
return
self._positive_video_option(CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD)
+ def read_video_codec_threads(self) -> int:
+ option = CoreOptions.READ_VIDEO_CODEC_THREADS
+ raw = self.options.to_map().get(option.key())
+ try:
+ value = self.options.get(option)
+ if isinstance(raw, (bool, float)) or value < 0:
+ raise ValueError()
+ except (TypeError, ValueError) as error:
+ raise ValueError(
+ "%s must be a non-negative integer." % option.key()) from error
+ return value
+
def _positive_video_option(self, option) -> int:
raw = self.options.to_map().get(option.key())
if isinstance(raw, (bool, float)):
diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
index 929b83d54d..8df8bd50e3 100644
--- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py
+++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py
@@ -285,13 +285,19 @@ class PaimonDatasetReader(ABC):
if isinstance(table_options, CoreOptions)
else CoreOptions.READ_VIDEO_DECODERS_PER_PAYLOAD.default_value()
)
+ codec_threads = (
+ table_options.read_video_codec_threads()
+ if isinstance(table_options, CoreOptions)
+ else CoreOptions.READ_VIDEO_CODEC_THREADS.default_value()
+ )
self._video_decode_pool = _VideoDecodePool()
self._video_collators = [
VideoFrameCollator(
access,
video_column=key,
decoder_factory=partial(
- _open_video_decoder, backend=self.video_backend),
+ _open_video_decoder, backend=self.video_backend,
+ codec_threads=codec_threads),
decode_batch_fn=_decode_video_frames,
output_column=key,
collate_fn=_identity,
@@ -1742,13 +1748,14 @@ def _video_decoder_fallback_errors():
return errors if ffmpeg_error is None else errors + (ffmpeg_error,)
-def _open_video_decoder(stream, backend=None):
+def _open_video_decoder(stream, backend=None, codec_threads=0):
+ pyav_options = {"thread_count": codec_threads} if codec_threads else {}
keyframe_index = getattr(stream, "video_keyframe_index", None)
if not isinstance(keyframe_index, VideoKeyframeIndex):
keyframe_index = None
if backend is None and keyframe_index is not None:
try:
- decoder = _PyAVVideoDecoder(stream, keyframe_index)
+ decoder = _PyAVVideoDecoder(stream, keyframe_index, **pyav_options)
except _video_decoder_fallback_errors():
stream.seek(0)
else:
@@ -1760,8 +1767,8 @@ def _open_video_decoder(stream, backend=None):
if backend == "torchcodec":
raise
stream.seek(0)
- return _PyAVVideoDecoder(stream) if keyframe_index is None \
- else _PyAVVideoDecoder(stream, keyframe_index)
+ return _PyAVVideoDecoder(stream, **pyav_options) if keyframe_index is None
\
+ else _PyAVVideoDecoder(stream, keyframe_index, **pyav_options)
def _open_torchcodec_decoder(stream):
@@ -1958,7 +1965,7 @@ class _PyAVVideoDecoder:
# Reuse common overlapping delta windows without retaining a whole video.
_FRAME_CACHE_SIZE = 8
- def __init__(self, stream, keyframe_index=None):
+ def __init__(self, stream, keyframe_index=None, thread_count=0):
try:
import av
except ImportError as error:
@@ -1968,6 +1975,7 @@ class _PyAVVideoDecoder:
) from error
self._source = stream
self._keyframe_index = keyframe_index
+ self._thread_count = thread_count
self._anchor_ordinals = (
[entry[0] for entry in keyframe_index.keyframes]
if keyframe_index is not None else [])
@@ -1981,6 +1989,7 @@ class _PyAVVideoDecoder:
if keyframe_index is None:
self._container = av.open(stream)
self._stream = self._container.streams.video[0]
+ self._stream.thread_count = self._thread_count
self._frames = iter(self._container.decode(self._stream))
def __getitem__(self, index):
@@ -2076,6 +2085,7 @@ class _PyAVVideoDecoder:
source.prefetch(self._keyframe_index.metadata_ranges)
with av.open(source) as container:
stream = container.streams.video[0]
+ stream.thread_count = self._thread_count
for anchor, targets in sorted(groups.items()):
unused_groups, ranges = self._indexed_plan(targets)
source.prefetch(ranges)
diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
index b718ca9d7f..331a7ef57c 100644
--- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
+++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py
@@ -730,7 +730,8 @@ class LeRobotValidationTest(unittest.TestCase):
for configured, expected in ((None, 16), ("2", 2)):
with self.subTest(configured=configured):
- options = {} if configured is None else {option: configured}
+ options = {} if configured is None else {
+ option: configured, "read.video.codec-threads": configured}
frames = SimpleNamespace(
options=CoreOptions(Options(options)))
rows = SimpleNamespace(
@@ -745,6 +746,9 @@ class LeRobotValidationTest(unittest.TestCase):
try:
self.assertEqual([expected, expected], [
c.max_open_videos for c in
dataset.reader._video_collators])
+ self.assertEqual([1 if configured is None else 2] * 2, [
+ c.decoder_factory.keywords["codec_threads"]
+ for c in dataset.reader._video_collators])
finally:
dataset.reader.close()
diff --git a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
index dfeb41bbb5..05a9025afd 100644
--- a/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
+++ b/paimon-python/pypaimon/tests/video_decode_parallelism_test.py
@@ -434,6 +434,68 @@ class VideoDecodeParallelismTest(unittest.TestCase):
torch.set_default_device(original_device)
collator.close()
+ def test_codec_thread_option(self):
+ self.assertEqual(1,
CoreOptions(Options({})).read_video_codec_threads())
+ for value in (0, 1, 4, "1"):
+ self.assertEqual(int(value), CoreOptions(Options({
+ "read.video.codec-threads":
value})).read_video_codec_threads())
+ for value in (-1, True, 1.5, "bad"):
+ with self.subTest(value=value), self.assertRaises(ValueError):
+ CoreOptions(Options({"read.video.codec-threads": value})) \
+ .read_video_codec_threads()
+
+ def test_codec_threads_preserve_real_indexed_and_scan_frames(self):
+ try:
+ import av
+ import numpy as np
+ import torch
+ except ImportError as error:
+ self.skipTest(str(error))
+ from pypaimon.multimodal.lerobot.dataset import _open_video_decoder
+ from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex
+
+ 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"
+ stream.codec_context.gop_size = 4
+ for value in range(12):
+ image = np.full((16, 16, 3), value * 15, 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)
+ payload = target.getvalue()
+ index = VideoKeyframeIndex.inspect(io.BytesIO(payload), len(payload))
+ outputs = []
+ for indexed in (False, True):
+ for threads in (0, 1, 2):
+ source = io.BytesIO(payload)
+ if indexed:
+ source.video_keyframe_index = index
+ source.video_length = len(payload)
+ decoder = _open_video_decoder(source, backend="pyav",
codec_threads=threads)
+ try:
+ if indexed:
+ decode_group = decoder._decode_indexed_group
+
+ def checked_group(container, stream, *args):
+ self.assertEqual(threads, stream.thread_count)
+ return decode_group(container, stream, *args)
+
+ with patch.object(decoder, '_decode_indexed_group',
checked_group):
+ frames = [decoder[i] for i in (9, 1, 7)]
+ else:
+ self.assertEqual(threads, decoder._stream.thread_count)
+ frames = [decoder[i] for i in (9, 1, 7)]
+ outputs.append(torch.stack(frames))
+ finally:
+ decoder.close()
+ source.close()
+ for output in outputs[1:]:
+ self.assertTrue(torch.equal(outputs[0], output))
+
def test_real_pyav_payloads_match_serial(self):
try:
import av
