From 4fad479a4035c8dd3652808f8911135c7101e4a6 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 26 Oct 2026 11:37:30 -0700 Subject: [PATCH 1/7] [python] Use persisted video indexes for TorchCodec frame reads --- docs/docs/pypaimon/lerobot.md | 8 +- .../pypaimon/multimodal/lerobot/dataset.py | 3 + .../multimodal/lerobot/indexed_torchcodec.py | 285 ++++++++++++++++++ .../pypaimon/tests/indexed_torchcodec_test.py | 276 +++++++++++++++++ 4 files changed, 569 insertions(+), 3 deletions(-) create mode 100644 paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py create mode 100644 paimon-python/pypaimon/tests/indexed_torchcodec_test.py diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 82c81590c733..f672115d3119 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -249,7 +249,9 @@ batches, avoiding a second full-resolution copy during collation. Without `tag_name`, the latest snapshots are used. Frame lookups use the BTree on `index`; payloads remain lazy. Indexed videos prefetch metadata and target GOPs, then fetch uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan path. Set -`video_backend` to force either decoder. +`video_backend` to force either decoder. With `video_backend="torchcodec"`, indexed +MP4 videos use timestamp seeks verified against their sample table and returned PTS. +Unsupported timelines or mismatched timestamps fall back to TorchCodec exact mode. With `pypaimon-rust` installed, LeRobot defaults to Rust frame lookup and table reads unless native read/planning options are explicitly configured. Set @@ -262,9 +264,9 @@ 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), +`read.video.codec-threads` controls FFmpeg threads per PyAV or indexed TorchCodec 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. +Set it to 0 for automatic selection. TorchCodec exact-mode fallback uses its own defaults. `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; diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index e1f55123600e..db180a953d31 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -1789,6 +1789,9 @@ def _open_video_decoder(stream, backend=None, codec_threads=0): stream.seek(0) else: return _FallbackVideoDecoder(stream, decoder) + if backend == "torchcodec" and keyframe_index is not None: + from pypaimon.multimodal.lerobot.indexed_torchcodec import _IndexedTorchCodecDecoder + return _IndexedTorchCodecDecoder(stream, keyframe_index, codec_threads) if backend in (None, "torchcodec"): try: return _open_torchcodec_decoder(stream) diff --git a/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py new file mode 100644 index 000000000000..e39fbe75bfbf --- /dev/null +++ b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py @@ -0,0 +1,285 @@ +# 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. + +"""Indexed TorchCodec reads with verified MP4 presentation timestamps.""" + +import bisect +import operator +import struct +from fractions import Fraction + +import numpy as np + + +_MAX_METADATA_BYTES = 16 * 1024 * 1024 +_MAX_SAMPLES = 1000000 +_MAX_TOP_LEVEL_BOXES = 1024 + + +def _boxes(data): + data = memoryview(data) + offset = 0 + box_count = 0 + while offset < len(data): + box_count += 1 + if box_count > 65536: + raise ValueError("Too many MP4 boxes.") + if len(data) - offset < 8: + raise ValueError("Truncated MP4 box.") + size, kind = struct.unpack_from(">I4s", data, offset) + header = 8 + if size == 1: + if len(data) - offset < 16: + raise ValueError("Truncated extended MP4 box.") + size = struct.unpack_from(">Q", data, offset + 8)[0] + header = 16 + elif size == 0: + size = len(data) - offset + if size < header or size > len(data) - offset: + raise ValueError("Invalid MP4 box size.") + yield kind, data[offset + header:offset + size] + offset += size + + +def _child(data, kind, required=True): + matches = [body for name, body in _boxes(data) if name == kind] + if len(matches) > 1 or (required and not matches): + raise ValueError("Missing or duplicate MP4 box: %r." % kind) + return matches[0] if matches else None + + +def _uint(data, offset): + if offset < 0 or offset + 4 > len(data): + raise ValueError("Truncated MP4 field.") + return struct.unpack_from(">I", data, offset)[0] + + +def _scale(data): + if not len(data) or data[0] not in (0, 1): + raise ValueError("Unsupported MP4 timing version.") + value = _uint(data, 20 if data[0] else 12) + if value == 0: + raise ValueError("Invalid MP4 time scale.") + return value + + +def _runs(data, composition=False): + count = _uint(data, 4) + if len(data) != 8 + count * 8 or count > _MAX_SAMPLES: + raise ValueError("Invalid MP4 timing run count.") + if data[0] not in ((0, 1) if composition else (0,)): + raise ValueError("Unsupported MP4 timing version.") + entries = np.frombuffer(data, dtype=">u4", offset=8).reshape(count, 2) + counts = entries[:, 0].astype(np.int64) + total = int(counts.sum()) + if total == 0 or total > _MAX_SAMPLES or np.any(counts == 0): + raise ValueError("MP4 sample count exceeds supported bounds.") + values = entries[:, 1] + if composition and data[0] == 1: + values = values.view(">i4") + return counts, values.astype(np.int64) + + +def _timeline(source, length): + source.seek(0) + moov = None + for unused in range(_MAX_TOP_LEVEL_BOXES): + offset = source.tell() + if offset >= length: + break + header = source.read(8) + if len(header) != 8: + raise ValueError("Truncated video header.") + size, kind = struct.unpack(">I4s", header) + header_size = 8 + if size == 1: + extended = source.read(8) + if len(extended) != 8: + raise ValueError("Truncated extended video header.") + size = struct.unpack(">Q", extended)[0] + header_size = 16 + elif size == 0: + size = length - offset + if size < header_size or size > length - offset: + raise ValueError("Invalid video box extent.") + if kind == b"moov": + if size > _MAX_METADATA_BYTES: + raise ValueError("MP4 metadata exceeds supported bounds.") + moov = source.read(size - header_size) + if len(moov) != size - header_size: + raise ValueError("Truncated MP4 metadata.") + break + source.seek(offset + size) + if moov is None or _child(moov, b"mvex", required=False) is not None: + raise ValueError("A non-fragmented MP4 sample table is required.") + tracks = [body for kind, body in _boxes(moov) if kind == b"trak"] + selected = None + for stream_index, track in enumerate(tracks): + media = _child(track, b"mdia") + handler = _child(media, b"hdlr") + if len(handler) < 12: + raise ValueError("Truncated MP4 handler.") + if bytes(handler[8:12]) == b"vide": + if selected is not None: + # The persisted index describes the first video track; keep + # TorchCodec's existing selection for multi-video containers. + raise ValueError("Multiple video tracks require exact mode.") + selected = (stream_index, track, media) + if selected is None: + raise ValueError("No video track.") + stream_index, track, media = selected + scale = _scale(_child(media, b"mdhd")) + samples = _child(_child(media, b"minf"), b"stbl") + counts, values = _runs(_child(samples, b"stts")) + count = int(counts.sum()) + if np.any(values <= 0) or _uint(_child(samples, b"stsz"), 8) != count: + raise ValueError("Invalid MP4 sample durations or count.") + durations = np.repeat(values, counts) + pts = np.cumsum(durations) - durations + composition = _child(samples, b"ctts", required=False) + if composition is not None: + counts, values = _runs(composition, composition=True) + if int(counts.sum()) != count: + raise ValueError("MP4 composition count mismatch.") + pts += np.repeat(values, counts) + order = np.argsort(pts, kind="stable") + pts, durations = pts[order], durations[order] + edits = _child(track, b"edts", required=False) + if edits is not None: + edit = _child(edits, b"elst") + n = _uint(edit, 4) + if edit[0] not in (0, 1) or n not in (1, 2): + raise ValueError("Unsupported MP4 edit list.") + fmt = struct.Struct(">Qqhh" if edit[0] else ">Iihh") + if len(edit) != 8 + n * fmt.size: + raise ValueError("Truncated MP4 edit list.") + entries = [fmt.unpack_from(edit, 8 + i * fmt.size) for i in range(n)] + if any((rate, fraction) != (1, 0) for _, _, rate, fraction in entries): + raise ValueError("Unsupported MP4 playback rate.") + if entries[-1][1] < 0 or (n == 2 and entries[0][1] != -1): + raise ValueError("Unsupported repeated MP4 media edit.") + movie_scale = _scale(_child(moov, b"mvhd")) + empty = entries[0][0] if n == 2 else 0 + shift = Fraction(empty * scale, movie_scale) - entries[-1][1] + limit = np.iinfo(np.int64) + if (shift.denominator != 1 or int(pts[0]) + shift < limit.min + or int(pts[-1]) + shift > limit.max): + raise ValueError("Unsupported MP4 edit shift.") + pts += int(shift) + start = Fraction(empty * scale, movie_scale) + end = start + Fraction(entries[-1][0] * scale, movie_scale) + # Cropping/discarded samples would change presentation ordinals. + if pts[0] < start or pts[-1] >= end: + raise ValueError("MP4 edit list removes samples.") + if pts[0] < 0 or np.any(pts[1:] <= pts[:-1]): + raise ValueError("Ambiguous MP4 presentation timestamps.") + source.seek(0) + return pts, durations, scale, stream_index + + +class _IndexedTorchCodecDecoder: + """Use timestamp seeks only when both the timeline and output are verified.""" + + def __init__(self, stream, index, codec_threads=1): + from pypaimon.multimodal.lerobot.dataset import ( + _RangeBackedVideo, _open_torchcodec_decoder) + from torchcodec.decoders import VideoDecoder + + self._stream = stream + self._index = index + self._range_source = None + self._decoder = None + self._indexed = False + self._exact = _open_torchcodec_decoder + self._ordinals = [entry[0] for entry in index.keyframes] + try: + self._length = operator.index(stream.video_length) + self._range_source = _RangeBackedVideo( + self._length, stream.video_read_ranges) + source = self._range_source + first = index.keyframes[0][2] + source.prefetch([(0, min(65536, self._length)), + (first, min(65536, self._length - first))] + + list(index.metadata_ranges)) + self._pts, self._durations, self._scale, stream_index = _timeline( + source, self._length) + if self._ordinals[-1] >= len(self._pts): + raise ValueError("Keyframe ordinal exceeds the MP4 timeline.") + # MP4 stream time base is 1 / media timescale. + if any(self._pts[o] != pts for o, pts, _ in index.keyframes): + raise ValueError("Keyframe PTS does not match the MP4 timeline.") + gaps = np.diff(self._pts / self._scale) + if len(gaps) and np.any(gaps <= 0): + raise ValueError("MP4 timestamps cannot be distinguished as seconds.") + self._pts_tolerance = min(1e-7, float(gaps.min()) / 4) if len(gaps) else 1e-7 + self._decoder = VideoDecoder( + source, stream_index=stream_index, device="cpu", + num_ffmpeg_threads=codec_threads, seek_mode="approximate") + self._indexed = True + except (AttributeError, TypeError, ValueError, OSError, RuntimeError, OverflowError): + self._fallback() + + def _fallback(self): + self._decoder = None + if self._range_source is not None: + self._range_source.close() + self._range_source = None + self._stream.seek(0) + self._decoder = self._exact(self._stream) + self._indexed = False + self._pts = self._durations = None + + def get_frames_at(self, *, indices): + positions = [operator.index(index) for index in indices] + if not self._indexed: + return self._decoder.get_frames_at(indices=positions) + if any(index < 0 or index >= len(self._pts) for index in positions): + raise IndexError("Video frame index is out of range.") + if not positions: + return self._decoder.get_frames_at(indices=positions) + ranges = [] + for index in sorted(set(positions)): + anchor = bisect.bisect_right(self._ordinals, index) - 1 + start = self._index.keyframes[anchor][2] + end = (self._index.keyframes[anchor + 1][2] + if anchor + 1 < len(self._ordinals) else self._length) + ranges.append((start, end - start)) + try: + self._range_source.prefetch(ranges) + expected = self._pts[positions] / self._scale + # Query inside the frame's presentation interval, avoiding floating + # point boundary rounding. Never return a neighbouring frame. + seconds = expected + self._durations[positions] / (2 * self._scale) + result = self._decoder.get_frames_played_at(seconds=seconds.tolist()) + actual = result.pts_seconds.cpu().numpy() + if actual.shape != expected.shape or not np.allclose( + actual, expected, rtol=0, atol=self._pts_tolerance): + raise ValueError("TorchCodec returned a different frame PTS.") + return result + except (TypeError, ValueError, OSError, RuntimeError, AttributeError, IndexError): + self._fallback() + return self._decoder.get_frames_at(indices=positions) + + def __getitem__(self, index): + return self.get_frames_at(indices=[index]).data[0] + + def close(self): + self._decoder = None + if self._range_source is not None: + self._range_source.close() + self._range_source = None + self._pts = self._durations = None diff --git a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py new file mode 100644 index 000000000000..38a9a8bbd0d3 --- /dev/null +++ b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py @@ -0,0 +1,276 @@ +# 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 struct +from fractions import Fraction +from unittest.mock import patch + +import numpy as np +import pytest + +from pypaimon.multimodal.lerobot.indexed_torchcodec import ( + _IndexedTorchCodecDecoder, _boxes, _runs, _timeline) +from pypaimon.table.row.video_keyframe_index import VideoKeyframeIndex + + +@pytest.fixture(params=[False, True]) +def mp4(request): + av = pytest.importorskip("av") + output = io.BytesIO() + with av.open(output, "w", format="mp4") as container: + stream = container.add_stream("libx264", rate=30) + stream.width, stream.height = 64, 48 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 12 + stream.codec_context.max_b_frames = 2 + pts = 0 + for i in range(60): + pixels = np.random.RandomState(i).randint(0, 256, (48, 64, 3), dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") + frame.pts, frame.time_base = pts, Fraction(1, 30) + pts += 1 + (i % 3 == 0) if request.param else 1 + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + return output.getvalue() + + +def indexed_source(payload): + source = io.BytesIO(payload) + source.video_length = len(payload) + source.video_keyframe_index = VideoKeyframeIndex.inspect(source, len(payload)) + source.seek(0) + calls = [] + + def read_ranges(ranges): + calls.extend(ranges) + return [payload[offset:offset + length] for offset, length in ranges] + + source.video_read_ranges = read_ranges + return source, calls + + +def require_torchcodec(): + try: + from torchcodec.decoders import VideoDecoder + except (ImportError, RuntimeError) as error: + pytest.skip("TorchCodec unavailable: %s" % error) + return VideoDecoder + + +def test_mp4_timeline_matches_full_decode(mp4): + import av + pts, durations, scale, stream = _timeline(io.BytesIO(mp4), len(mp4)) + with av.open(io.BytesIO(mp4)) as container: + expected = [float(frame.pts * frame.time_base) for frame in container.decode(video=0)] + np.testing.assert_allclose(pts / scale, expected, rtol=0, atol=1e-9) + assert len(durations) == len(expected) + assert stream == 0 + + +def test_indexed_torchcodec_matches_exact_frames(mp4): + VideoDecoder = require_torchcodec() + source, calls = indexed_source(mp4) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + assert decoder._indexed + indices = [59, 0, 24, 24, 13, 35, 1] + expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=indices) + actual = decoder.get_frames_at(indices=indices) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) + assert calls + # CFR stays indexed. TorchCodec may reject VFR timestamp seeks; + # those must still return the exact requested frames through fallback. + pts, unused, unused, unused = _timeline(io.BytesIO(mp4), len(mp4)) + if len(np.unique(np.diff(pts))) == 1: + assert decoder._indexed + with pytest.raises(IndexError): + decoder.get_frames_at(indices=[60]) + finally: + decoder.close() + assert not source.closed + + +def test_wrong_returned_pts_falls_back_to_exact(mp4): + VideoDecoder = require_torchcodec() + source, unused = indexed_source(mp4) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[20]) + wrong = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[21]) + try: + with patch.object(decoder._decoder, "get_frames_played_at", return_value=wrong): + actual = decoder.get_frames_at(indices=[20]) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + assert not decoder._indexed + finally: + decoder.close() + + +def test_unsupported_timeline_falls_back_to_exact(mp4): + VideoDecoder = require_torchcodec() + source, unused = indexed_source(mp4) + with patch("pypaimon.multimodal.lerobot.indexed_torchcodec._timeline", + side_effect=ValueError("unsupported edit")): + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + assert not decoder._indexed + actual = decoder[7] + expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact")[7] + np.testing.assert_array_equal(actual.numpy(), expected.numpy()) + finally: + decoder.close() + + +def test_timing_limits_before_expanding(): + data = struct.pack(">IIII", 0, 1, 1000001, 1) + with pytest.raises(ValueError, match="bounds"): + _runs(data) + with pytest.raises(ValueError, match="count"): + _runs(data[:-1]) + + +@pytest.mark.parametrize("data", [b"x", struct.pack(">I4s", 4, b"moov"), + struct.pack(">I4s", 1, b"moov"), + struct.pack(">I4s", 100, b"moov")]) +def test_malformed_box_is_rejected(data): + with pytest.raises(ValueError): + list(_boxes(data)) + + +def test_metadata_limit_before_reading_body(): + class Source(io.BytesIO): + def read(self, size=-1): + assert size <= 8 + return super().read(size) + + size = 17 * 1024 * 1024 + source = Source(struct.pack(">I4s", size, b"moov")) + with pytest.raises(ValueError, match="bounds"): + _timeline(source, size) + + +def test_first_decode_error_falls_back(mp4): + VideoDecoder = require_torchcodec() + source, unused = indexed_source(mp4) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + with patch.object(decoder._decoder, "get_frames_played_at", + side_effect=IndexError("approximate timestamp")): + actual = decoder.get_frames_at(indices=[5, 0, 5]) + expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[5, 0, 5]) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + assert not decoder._indexed + finally: + decoder.close() + + +def test_explicit_backend_routes_indexed_videos(mp4): + require_torchcodec() + from pypaimon.multimodal.lerobot.dataset import _open_video_decoder + source, unused = indexed_source(mp4) + decoder = _open_video_decoder(source, backend="torchcodec", codec_threads=1) + try: + assert isinstance(decoder, _IndexedTorchCodecDecoder) + assert decoder._indexed + finally: + decoder.close() + + +@pytest.mark.parametrize("fragmented,audio_first", [(True, False), (False, True)]) +def test_fragmented_and_audio_first_video(fragmented, audio_first): + av = pytest.importorskip("av") + VideoDecoder = require_torchcodec() + output = io.BytesIO() + options = {"movflags": "frag_keyframe+empty_moov+default_base_moof"} if fragmented else None + with av.open(output, "w", format="mp4", options=options) as container: + if audio_first: + audio = container.add_stream("aac", rate=48000) + audio.layout = "mono" + video = container.add_stream("libx264", rate=30) + video.width, video.height = 64, 48 + video.pix_fmt = "yuv420p" + video.codec_context.gop_size = 12 + video.codec_context.max_b_frames = 2 + if audio_first: + for i in range(30): + frame = av.AudioFrame.from_ndarray(np.zeros((1, 1600), dtype=np.float32), + format="fltp", layout="mono") + frame.sample_rate = 48000 + frame.pts = i * 1600 + frame.time_base = Fraction(1, 48000) + for packet in audio.encode(frame): + container.mux(packet) + for packet in audio.encode(): + container.mux(packet) + for i in range(30): + frame = av.VideoFrame.from_ndarray(np.full((48, 64, 3), i * 7, dtype=np.uint8), format="rgb24") + for packet in video.encode(frame): + container.mux(packet) + for packet in video.encode(): + container.mux(packet) + payload = output.getvalue() + source, unused = indexed_source(payload) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + assert decoder._indexed is (not fragmented) + positions = [29, 0, 16, 4] + actual = decoder.get_frames_at(indices=positions) + expected = VideoDecoder(io.BytesIO(payload), stream_index=1 if audio_first else 0, + seek_mode="exact").get_frames_at(indices=positions) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) + finally: + decoder.close() + + +def test_multiple_video_tracks_preserve_exact_selection(): + av = pytest.importorskip("av") + VideoDecoder = require_torchcodec() + output = io.BytesIO() + with av.open(output, "w", format="mp4") as container: + streams = [container.add_stream("libx264", rate=30) for _ in range(2)] + for i, stream in enumerate(streams): + stream.width, stream.height = 32 * (i + 1), 32 * (i + 1) + stream.pix_fmt = "yuv420p" + for i in range(12): + for stream in streams: + pixels = np.full((stream.height, stream.width, 3), i * 16, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for stream in streams: + for packet in stream.encode(): + container.mux(packet) + payload = output.getvalue() + source, unused = indexed_source(payload) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + assert not decoder._indexed + actual = decoder.get_frames_at(indices=[11, 0, 5]) + expected = VideoDecoder(io.BytesIO(payload), seek_mode="exact").get_frames_at(indices=[11, 0, 5]) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + finally: + decoder.close() + + +def test_signed_composition_offsets(): + counts, values = _runs(struct.pack(">IIIi", 1 << 24, 1, 3, -2), composition=True) + np.testing.assert_array_equal(counts, [3]) + np.testing.assert_array_equal(values, [-2]) From cf35c4228b7dfacd3ebe22e0f7b362ca9d9a3d1f Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Mon, 26 Oct 2026 11:41:31 -0700 Subject: [PATCH 2/7] [python] Clarify codec thread option for indexed TorchCodec --- paimon-python/pypaimon/common/options/core_options.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index d462943b93f9..78ced0fbe4f1 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -1185,7 +1185,7 @@ class CoreOptions: .int_type() .default_value(1) .with_description( - "FFmpeg threads per PyAV video decoder in a LeRobot worker. " + "FFmpeg threads per PyAV or indexed TorchCodec decoder in a LeRobot worker. " "0 keeps automatic selection; 1 avoids nested codec parallelism.") ) From 971818952b1194cf4c6b50ac959c3e41ccab877f Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Fri, 30 Oct 2026 08:44:37 -0700 Subject: [PATCH 3/7] [python] Share compact TorchCodec timelines across decoder replicas --- docs/docs/pypaimon/lerobot.md | 1 + .../multimodal/lerobot/indexed_torchcodec.py | 88 ++++++++++--- paimon-python/pypaimon/multimodal/video.py | 7 ++ .../pypaimon/tests/indexed_torchcodec_test.py | 116 +++++++++++++++++- 4 files changed, 195 insertions(+), 17 deletions(-) diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index f672115d3119..8726868efcd7 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -251,6 +251,7 @@ on `index`; payloads remain lazy. Indexed videos prefetch metadata and target GO uncached PyAV reads on demand. Unindexed videos use the TorchCodec/PyAV scan path. Set `video_backend` to force either decoder. With `video_backend="torchcodec"`, indexed MP4 videos use timestamp seeks verified against their sample table and returned PTS. +Decoder replicas share read-only timelines; uniform timing is stored compactly. Unsupported timelines or mismatched timestamps fall back to TorchCodec exact mode. With `pypaimon-rust` installed, LeRobot defaults to Rust frame lookup and table reads diff --git a/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py index e39fbe75bfbf..3cbc1c953adf 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py +++ b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py @@ -30,6 +30,56 @@ _MAX_TOP_LEVEL_BOXES = 1024 +class _TimingValues: + """Keep arithmetic sequences compact, otherwise retain a read-only array.""" + + def __init__(self, values): + self.count = len(values) + self.start = int(values[0]) + self.step = int(values[1]) - self.start if self.count > 1 else 0 + self.values = None + for start in range(1, self.count, 65536): + end = min(start + 65536, self.count) + if np.any(values[start:end] - values[start - 1:end - 1] != self.step): + values.flags.writeable = False + self.values = values + break + + def __len__(self): + return self.count + + def __getitem__(self, positions): + if self.values is not None: + return self.values[positions] + if isinstance(positions, (int, np.integer)): + if positions < 0: + positions += self.count + if not 0 <= positions < self.count: + raise IndexError("Timing position is out of range.") + return self.start + int(positions) * self.step + positions = np.asarray(positions, dtype=np.int64) + return self.start + positions * self.step + + +class _SharedTimeline: + def __init__(self, source, length): + pts, durations, self.scale, self.stream_index = _timeline(source, length) + minimum = float('inf') + previous = None + for start in range(0, len(pts), 65536): + seconds = pts[start:start + 65536] / self.scale + if previous is not None: + minimum = min(minimum, float(seconds[0] - previous)) + if len(seconds) > 1: + minimum = min(minimum, float(np.diff(seconds).min())) + previous = seconds[-1] + if minimum <= 0: + raise ValueError("MP4 timestamps cannot be distinguished as seconds.") + self.tolerance = min(1e-7, minimum / 4) + self.pts = _TimingValues(pts) + self.durations = _TimingValues(durations) + + def _boxes(data): data = memoryview(data) offset = 0 @@ -149,15 +199,17 @@ def _timeline(source, length): if np.any(values <= 0) or _uint(_child(samples, b"stsz"), 8) != count: raise ValueError("Invalid MP4 sample durations or count.") durations = np.repeat(values, counts) - pts = np.cumsum(durations) - durations + pts = np.cumsum(durations) + pts -= durations composition = _child(samples, b"ctts", required=False) if composition is not None: counts, values = _runs(composition, composition=True) if int(counts.sum()) != count: raise ValueError("MP4 composition count mismatch.") pts += np.repeat(values, counts) - order = np.argsort(pts, kind="stable") - pts, durations = pts[order], durations[order] + if np.any(pts[1:] <= pts[:-1]): + order = np.argsort(pts, kind="stable") + pts, durations = pts[order], durations[order] edits = _child(track, b"edts", required=False) if edits is not None: edit = _child(edits, b"elst") @@ -205,7 +257,7 @@ def __init__(self, stream, index, codec_threads=1): self._decoder = None self._indexed = False self._exact = _open_torchcodec_decoder - self._ordinals = [entry[0] for entry in index.keyframes] + self._timeline = None try: self._length = operator.index(stream.video_length) self._range_source = _RangeBackedVideo( @@ -215,19 +267,23 @@ def __init__(self, stream, index, codec_threads=1): source.prefetch([(0, min(65536, self._length)), (first, min(65536, self._length - first))] + list(index.metadata_ranges)) - self._pts, self._durations, self._scale, stream_index = _timeline( - source, self._length) - if self._ordinals[-1] >= len(self._pts): + cache = getattr(stream, "video_metadata_cache", None) + cache_key = getattr(stream, "video_metadata_key", None) + timeline = cache.get(cache_key) if cache is not None else None + if timeline is None: + timeline = _SharedTimeline(source, self._length) + self._pts, self._durations, self._scale = timeline.pts, timeline.durations, timeline.scale + if index.keyframes[-1][0] >= len(self._pts): raise ValueError("Keyframe ordinal exceeds the MP4 timeline.") # MP4 stream time base is 1 / media timescale. if any(self._pts[o] != pts for o, pts, _ in index.keyframes): raise ValueError("Keyframe PTS does not match the MP4 timeline.") - gaps = np.diff(self._pts / self._scale) - if len(gaps) and np.any(gaps <= 0): - raise ValueError("MP4 timestamps cannot be distinguished as seconds.") - self._pts_tolerance = min(1e-7, float(gaps.min()) / 4) if len(gaps) else 1e-7 + self._timeline = timeline + self._pts_tolerance = timeline.tolerance + if cache is not None: + cache[cache_key] = timeline self._decoder = VideoDecoder( - source, stream_index=stream_index, device="cpu", + source, stream_index=timeline.stream_index, device="cpu", num_ffmpeg_threads=codec_threads, seek_mode="approximate") self._indexed = True except (AttributeError, TypeError, ValueError, OSError, RuntimeError, OverflowError): @@ -241,7 +297,7 @@ def _fallback(self): self._stream.seek(0) self._decoder = self._exact(self._stream) self._indexed = False - self._pts = self._durations = None + self._pts = self._durations = self._timeline = None def get_frames_at(self, *, indices): positions = [operator.index(index) for index in indices] @@ -253,10 +309,10 @@ def get_frames_at(self, *, indices): return self._decoder.get_frames_at(indices=positions) ranges = [] for index in sorted(set(positions)): - anchor = bisect.bisect_right(self._ordinals, index) - 1 + anchor = bisect.bisect_right(self._index.keyframes, (index, float("inf"), float("inf"))) - 1 start = self._index.keyframes[anchor][2] end = (self._index.keyframes[anchor + 1][2] - if anchor + 1 < len(self._ordinals) else self._length) + if anchor + 1 < len(self._index.keyframes) else self._length) ranges.append((start, end - start)) try: self._range_source.prefetch(ranges) @@ -282,4 +338,4 @@ def close(self): if self._range_source is not None: self._range_source.close() self._range_source = None - self._pts = self._durations = None + self._pts = self._durations = self._timeline = None diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index a45b01b0a0cf..6316b156e1ac 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -19,6 +19,7 @@ import io import os +import weakref from collections import OrderedDict, deque from collections.abc import Mapping from concurrent.futures import ThreadPoolExecutor, wait @@ -224,6 +225,7 @@ def __init__( self.collate_fn = collate_fn self.range_parallelism = range_parallelism self._decoders = OrderedDict() + self._video_metadata_cache = weakref.WeakValueDictionary() self._owner_pid = os.getpid() def __call__(self, rows): @@ -238,6 +240,7 @@ def __call__(self, rows): def close(self): resources = list(self._decoders.values()) self._decoders.clear() + self._video_metadata_cache = weakref.WeakValueDictionary() first_error = None for resource in resources: try: @@ -253,6 +256,7 @@ def __getstate__(self): # stream state opened in the parent process. state = self.__dict__.copy() state["_decoders"] = OrderedDict() + state["_video_metadata_cache"] = None state["_owner_pid"] = None return state @@ -382,6 +386,9 @@ def _decoder(self, descriptor, keyframe_index_descriptor=None, replica=0): descriptor.length, ).new_input_stream() try: + # Replica decoders share metadata, while their decode state stays separate. + stream.video_metadata_cache = self._video_metadata_cache + stream.video_metadata_key = (descriptor, keyframe_index_descriptor) if keyframe_index_descriptor is not None: mapping = self.file_io.read_file_range( keyframe_index_descriptor.uri, diff --git a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py index 38a9a8bbd0d3..de794adaa678 100644 --- a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py +++ b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py @@ -101,7 +101,9 @@ def test_indexed_torchcodec_matches_exact_frames(mp4): pts, unused, unused, unused = _timeline(io.BytesIO(mp4), len(mp4)) if len(np.unique(np.diff(pts))) == 1: assert decoder._indexed - with pytest.raises(IndexError): + # Exact-mode TorchCodec versions differ in their out-of-range exception. + errors = IndexError if decoder._indexed else (IndexError, RuntimeError) + with pytest.raises(errors, match="[Ii]nvalid frame index|out of range"): decoder.get_frames_at(indices=[60]) finally: decoder.close() @@ -274,3 +276,115 @@ def test_signed_composition_offsets(): counts, values = _runs(struct.pack(">IIIi", 1 << 24, 1, 3, -2), composition=True) np.testing.assert_array_equal(counts, [3]) np.testing.assert_array_equal(values, [-2]) + + +def timing_only_mp4(count, variable=False): + def box(kind, body): + return struct.pack('>I4s', len(body) + 8, kind) + body + + if variable: + runs = np.ones((count, 2), dtype='>u4') + runs[:, 1] = 1001 + runs[1::2, 1] = 1002 + stts = box(b'stts', struct.pack('>II', 0, count) + runs.tobytes()) + else: + stts = box(b'stts', struct.pack('>IIII', 0, 1, count, 1001)) + stsz = box(b'stsz', struct.pack('>III', 0, 1, count)) + media = box(b'mdhd', b'\0' * 12 + struct.pack('>I', 30000)) + media += box(b'hdlr', b'\0' * 8 + b'vide') + media += box(b'minf', box(b'stbl', stts + stsz)) + return box(b'moov', box(b'trak', box(b'mdia', media))) + + +@pytest.mark.parametrize("variable", [False, True]) +def test_million_sample_timeline_memory(variable): + import tracemalloc + from pypaimon.multimodal.lerobot.indexed_torchcodec import _MAX_SAMPLES, _SharedTimeline + + payload = timing_only_mp4(_MAX_SAMPLES, variable=variable) + tracemalloc.start() + try: + timeline = _SharedTimeline(io.BytesIO(payload), len(payload)) + retained, peak = tracemalloc.get_traced_memory() + finally: + tracemalloc.stop() + if variable: + assert timeline.pts.values.nbytes + timeline.durations.values.nbytes == 16_000_000 + assert retained < 17 * 1024 * 1024 + assert peak < 64 * 1024 * 1024 + else: + assert timeline.pts.values is None + assert timeline.durations.values is None + assert retained < 1024 * 1024 + assert peak < 24 * 1024 * 1024 + positions = np.array([0, 500000, _MAX_SAMPLES - 1, 0]) + np.testing.assert_array_equal(timeline.pts[positions], positions * 1001 + (positions // 2 if variable else 0)) + np.testing.assert_array_equal(timeline.durations[positions], 1001 + positions % 2 if variable else [1001] * 4) + + +def test_irregular_timing_arrays_are_read_only(): + from pypaimon.multimodal.lerobot.indexed_torchcodec import _TimingValues + + original = np.array([0, 1001, 3003, 4004], dtype=np.int64) + timing = _TimingValues(original) + assert timing.values is original + assert not timing.values.flags.writeable + np.testing.assert_array_equal(timing[[3, 0, 1]], [4004, 0, 1001]) + with pytest.raises(ValueError): + timing.values[0] = 1 + + +@pytest.mark.parametrize("variable", [False, True]) +def test_replica_decoders_share_and_release_timeline(tmp_path, variable): + import sys + import weakref + from types import SimpleNamespace + from unittest.mock import Mock + from pypaimon.common.file_io import FileIO + from pypaimon.multimodal.video import VideoFrameCollator + from pypaimon.table.row.blob import VideoFrameDescriptor + + payload = timing_only_mp4(1000000, variable=variable) + mapping = VideoKeyframeIndex([(0, len(payload))], [(0, 0, 0)]).serialize() + path = tmp_path / 'replicas.video' + path.write_bytes(payload + mapping) + frame = VideoFrameDescriptor(path.as_uri(), 0, len(payload), 0, len(payload), len(mapping)) + table = SimpleNamespace(file_io=FileIO.get(path.as_uri(), {})) + collator = VideoFrameCollator( + table, video_column='video', max_open_videos=4, + decoder_factory=lambda stream: _IndexedTorchCodecDecoder(stream, stream.video_keyframe_index), + decode_fn=lambda *args: None) + fake = SimpleNamespace(VideoDecoder=Mock()) + with patch.dict(sys.modules, {'torchcodec.decoders': fake}), patch( + 'pypaimon.multimodal.lerobot.indexed_torchcodec._timeline', wraps=_timeline) as parse: + try: + decoders = [collator._decoder(frame.payload_descriptor, frame.keyframe_index_descriptor, replica=i) + for i in range(4)] + assert all(decoder._indexed for decoder in decoders) + assert parse.call_count == 1 + assert all(decoder._timeline is decoders[0]._timeline for decoder in decoders) + reference = weakref.ref(decoders[0]._timeline) + state = collator.__getstate__() + assert state['_video_metadata_cache'] is None + worker = VideoFrameCollator.__new__(VideoFrameCollator) + worker.__dict__.update(state) + worker._ensure_process_local_cache() + assert len(worker._video_metadata_cache) == 0 + assert len(collator._video_metadata_cache) == 1 + worker.close() + decoders[0]._fallback() + assert reference() is not None + assert decoders[1]._pts[999999] == 999999 * 1001 + (499999 if variable else 0) + if variable: + assert decoders[1]._pts.values.nbytes == 8_000_000 + assert all(decoder._pts is decoders[1]._pts for decoder in decoders[1:]) + assert all(decoder._durations is decoders[1]._durations for decoder in decoders[1:]) + other_path = tmp_path / 'another.video' + other_path.write_bytes(payload + mapping) + other_frame = VideoFrameDescriptor(other_path.as_uri(), 0, len(payload), 0, len(payload), len(mapping)) + other = collator._decoder(other_frame.payload_descriptor, other_frame.keyframe_index_descriptor) + assert other._timeline is not reference() + assert parse.call_count == 2 + finally: + collator.close() + assert reference() is None From 302cf3b71662d0784e8356d984ef4649a964880e Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Sat, 31 Oct 2026 08:02:45 -0700 Subject: [PATCH 4/7] Require indexed TorchCodec tests in LeRobot compatibility CI --- .github/workflows/ci-python.yml | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 580a6a2a78a8..0ffdc1576024 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -302,6 +302,18 @@ jobs: pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_boolean_statistics_preserve_extrema \ pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_int64_statistics_preserve_extrema + - name: Test indexed TorchCodec (no skips) + working-directory: paimon-python + run: | + python -c 'import av; from torchcodec.decoders import VideoDecoder' + python -m pytest -q pypaimon/tests/indexed_torchcodec_test.py --junitxml=indexed-torchcodec.xml + python - <<'PYTEST_REPORT' + import xml.etree.ElementTree as ET + cases = list(ET.parse("indexed-torchcodec.xml").iter("testcase")) + assert cases, "No indexed TorchCodec tests ran" + assert not any(case.find("skipped") is not None for case in cases), "Indexed TorchCodec tests must not skip" + PYTEST_REPORT + torch: name: PyTorch / Python 3.10 timeout-minutes: 60 From 26ab105ccbab286fbac1930195fc0656ad301fd1 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Sat, 31 Oct 2026 08:51:20 -0700 Subject: [PATCH 5/7] Test TorchCodec with compatible FFmpeg and bytes-only fallback --- .github/workflows/ci-python.yml | 20 +++++-- .../pypaimon/tests/indexed_torchcodec_test.py | 54 ++++++++++++++----- 2 files changed, 59 insertions(+), 15 deletions(-) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 0ffdc1576024..58d22197cc5c 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -267,19 +267,33 @@ jobs: done lerobot-compatibility: - name: LeRobot ${{ matrix.lerobot }} / Python ${{ matrix.python }} + name: LeRobot ${{ matrix.lerobot }} / Python ${{ matrix.python }} / TorchCodec ${{ matrix.torchcodec }} timeout-minutes: 40 - runs-on: ubuntu-latest + runs-on: ubuntu-24.04 strategy: fail-fast: false matrix: include: - python: "3.10" lerobot: "0.4.4" + torch: "2.7.1" + torchvision: "0.22.1" + torchcodec: "0.3.0" + - python: "3.10" + lerobot: "0.4.4" + torch: "2.6.0" + torchvision: "0.21.0" + torchcodec: "0.2.1" - python: "3.12" lerobot: "0.6.0" + torch: "2.7.1" + torchvision: "0.22.1" + torchcodec: "0.3.0" - python: "3.12" lerobot: "0.6.1" + torch: "2.7.1" + torchvision: "0.22.1" + torchcodec: "0.3.0" steps: - uses: actions/checkout@v6 - uses: actions/setup-python@v6 @@ -290,7 +304,7 @@ jobs: - name: Install LeRobot dependencies run: | python -m pip install --upgrade pip - python -m pip install torch==2.7.1 torchvision==0.22.1 torchcodec==0.3.0 --index-url https://download.pytorch.org/whl/cpu + python -m pip install torch==${{ matrix.torch }} torchvision==${{ matrix.torchvision }} torchcodec==${{ matrix.torchcodec }} --index-url https://download.pytorch.org/whl/cpu python -m pip install -e './paimon-python[lerobot]' 'lerobot==${{ matrix.lerobot }}' --group ./paimon-python/pyproject.toml:test python -m pip check python -c 'from lerobot.datasets.lerobot_dataset import LeRobotDataset' diff --git a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py index de794adaa678..acdbd6101dc7 100644 --- a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py +++ b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py @@ -74,6 +74,12 @@ def require_torchcodec(): return VideoDecoder +def supports_file_like(): + import torchcodec + from packaging.version import Version + return Version(torchcodec.__version__) >= Version("0.3") + + def test_mp4_timeline_matches_full_decode(mp4): import av pts, durations, scale, stream = _timeline(io.BytesIO(mp4), len(mp4)) @@ -89,9 +95,9 @@ def test_indexed_torchcodec_matches_exact_frames(mp4): source, calls = indexed_source(mp4) decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) try: - assert decoder._indexed + assert decoder._indexed is supports_file_like() indices = [59, 0, 24, 24, 13, 35, 1] - expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=indices) + expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=indices) actual = decoder.get_frames_at(indices=indices) np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) @@ -99,8 +105,8 @@ def test_indexed_torchcodec_matches_exact_frames(mp4): # CFR stays indexed. TorchCodec may reject VFR timestamp seeks; # those must still return the exact requested frames through fallback. pts, unused, unused, unused = _timeline(io.BytesIO(mp4), len(mp4)) - if len(np.unique(np.diff(pts))) == 1: - assert decoder._indexed + if supports_file_like() and len(np.unique(np.diff(pts))) == 1: + assert decoder._indexed is supports_file_like() # Exact-mode TorchCodec versions differ in their out-of-range exception. errors = IndexError if decoder._indexed else (IndexError, RuntimeError) with pytest.raises(errors, match="[Ii]nvalid frame index|out of range"): @@ -114,8 +120,8 @@ def test_wrong_returned_pts_falls_back_to_exact(mp4): VideoDecoder = require_torchcodec() source, unused = indexed_source(mp4) decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) - expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[20]) - wrong = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[21]) + expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[20]) + wrong = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[21]) try: with patch.object(decoder._decoder, "get_frames_played_at", return_value=wrong): actual = decoder.get_frames_at(indices=[20]) @@ -125,6 +131,30 @@ def test_wrong_returned_pts_falls_back_to_exact(mp4): decoder.close() +def test_bytes_only_decoder_falls_back_to_exact(mp4): + VideoDecoder = require_torchcodec() + source, unused = indexed_source(mp4) + inputs = [] + + def bytes_only(data, **kwargs): + inputs.append(type(data)) + if not isinstance(data, bytes): + raise TypeError("file-like input unsupported") + return VideoDecoder(data, **kwargs) + + with patch("torchcodec.decoders.VideoDecoder", side_effect=bytes_only): + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + assert not decoder._indexed + assert inputs[-1] is bytes + actual = decoder.get_frames_at(indices=[59, 0, 20, 20]) + expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[59, 0, 20, 20]) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) + finally: + decoder.close() + + def test_unsupported_timeline_falls_back_to_exact(mp4): VideoDecoder = require_torchcodec() source, unused = indexed_source(mp4) @@ -134,7 +164,7 @@ def test_unsupported_timeline_falls_back_to_exact(mp4): try: assert not decoder._indexed actual = decoder[7] - expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact")[7] + expected = VideoDecoder(mp4, seek_mode="exact")[7] np.testing.assert_array_equal(actual.numpy(), expected.numpy()) finally: decoder.close() @@ -176,7 +206,7 @@ def test_first_decode_error_falls_back(mp4): with patch.object(decoder._decoder, "get_frames_played_at", side_effect=IndexError("approximate timestamp")): actual = decoder.get_frames_at(indices=[5, 0, 5]) - expected = VideoDecoder(io.BytesIO(mp4), seek_mode="exact").get_frames_at(indices=[5, 0, 5]) + expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[5, 0, 5]) np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) assert not decoder._indexed finally: @@ -190,7 +220,7 @@ def test_explicit_backend_routes_indexed_videos(mp4): decoder = _open_video_decoder(source, backend="torchcodec", codec_threads=1) try: assert isinstance(decoder, _IndexedTorchCodecDecoder) - assert decoder._indexed + assert decoder._indexed is supports_file_like() finally: decoder.close() @@ -231,10 +261,10 @@ def test_fragmented_and_audio_first_video(fragmented, audio_first): source, unused = indexed_source(payload) decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) try: - assert decoder._indexed is (not fragmented) + assert decoder._indexed is (supports_file_like() and not fragmented) positions = [29, 0, 16, 4] actual = decoder.get_frames_at(indices=positions) - expected = VideoDecoder(io.BytesIO(payload), stream_index=1 if audio_first else 0, + expected = VideoDecoder(payload, stream_index=1 if audio_first else 0, seek_mode="exact").get_frames_at(indices=positions) np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) @@ -266,7 +296,7 @@ def test_multiple_video_tracks_preserve_exact_selection(): try: assert not decoder._indexed actual = decoder.get_frames_at(indices=[11, 0, 5]) - expected = VideoDecoder(io.BytesIO(payload), seek_mode="exact").get_frames_at(indices=[11, 0, 5]) + expected = VideoDecoder(payload, seek_mode="exact").get_frames_at(indices=[11, 0, 5]) np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) finally: decoder.close() From 97e60f771373e854461fb6910d9da319fb8718a9 Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Sat, 31 Oct 2026 09:05:14 -0700 Subject: [PATCH 6/7] Limit TorchCodec 0.2 compatibility job to indexed decoder tests --- .github/workflows/ci-python.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 58d22197cc5c..87f7a9d05d3b 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -307,10 +307,11 @@ jobs: python -m pip install torch==${{ matrix.torch }} torchvision==${{ matrix.torchvision }} torchcodec==${{ matrix.torchcodec }} --index-url https://download.pytorch.org/whl/cpu python -m pip install -e './paimon-python[lerobot]' 'lerobot==${{ matrix.lerobot }}' --group ./paimon-python/pyproject.toml:test python -m pip check - python -c 'from lerobot.datasets.lerobot_dataset import LeRobotDataset' - name: Test import, training reads and writer statistics + if: matrix.torchcodec != '0.2.1' working-directory: paimon-python run: | + python -c 'from lerobot.datasets.lerobot_dataset import LeRobotDataset' python -m pytest -q \ pypaimon/tests/multimodal_lerobot_test.py \ pypaimon/tests/multimodal_lerobot_writer_test.py::PaimonLeRobotWriterTest::test_multiframe_boolean_statistics_preserve_extrema \ From 4add1beb07c9b6e6a7b019b4805042287277dbff Mon Sep 17 00:00:00 2001 From: "xiaohongbo.xhb" Date: Sat, 31 Oct 2026 09:07:08 -0700 Subject: [PATCH 7/7] Fix indexed TorchCodec seeks when nominal FPS differs from timestamps --- .../multimodal/lerobot/indexed_torchcodec.py | 15 ++++- .../pypaimon/tests/indexed_torchcodec_test.py | 55 ++++++++++++++++++- 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py index 3cbc1c953adf..6df7690e6a6f 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py +++ b/paimon-python/pypaimon/multimodal/lerobot/indexed_torchcodec.py @@ -320,7 +320,20 @@ def get_frames_at(self, *, indices): # Query inside the frame's presentation interval, avoiding floating # point boundary rounding. Never return a neighbouring frame. seconds = expected + self._durations[positions] / (2 * self._scale) - result = self._decoder.get_frames_played_at(seconds=seconds.tolist()) + # TorchCodec's approximate batch API may turn timestamps into + # ordinals using the nominal header frame rate. Remuxed videos can + # have a different sample-table cadence. The scalar timestamp API + # seeks by PTS instead; sort/deduplicate requests, then restore order. + import torch + from torchcodec import FrameBatch + frames = {position: self._decoder.get_frame_played_at(float(second)) + for position, second in sorted(set(zip(positions, seconds)))} + ordered = [frames[position] for position in positions] + result = FrameBatch( + data=torch.stack([frame.data for frame in ordered]), + pts_seconds=torch.tensor([frame.pts_seconds for frame in ordered], dtype=torch.float64), + duration_seconds=torch.tensor( + [frame.duration_seconds for frame in ordered], dtype=torch.float64)) actual = result.pts_seconds.cpu().numpy() if actual.shape != expected.shape or not np.allclose( actual, expected, rtol=0, atol=self._pts_tolerance): diff --git a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py index acdbd6101dc7..3c0cf939801e 100644 --- a/paimon-python/pypaimon/tests/indexed_torchcodec_test.py +++ b/paimon-python/pypaimon/tests/indexed_torchcodec_test.py @@ -121,9 +121,9 @@ def test_wrong_returned_pts_falls_back_to_exact(mp4): source, unused = indexed_source(mp4) decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[20]) - wrong = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[21]) + wrong = VideoDecoder(mp4, seek_mode="exact").get_frame_at(21) try: - with patch.object(decoder._decoder, "get_frames_played_at", return_value=wrong): + with patch.object(decoder._decoder, "get_frame_played_at", return_value=wrong): actual = decoder.get_frames_at(indices=[20]) np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) assert not decoder._indexed @@ -203,7 +203,7 @@ def test_first_decode_error_falls_back(mp4): source, unused = indexed_source(mp4) decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) try: - with patch.object(decoder._decoder, "get_frames_played_at", + with patch.object(decoder._decoder, "get_frame_played_at", side_effect=IndexError("approximate timestamp")): actual = decoder.get_frames_at(indices=[5, 0, 5]) expected = VideoDecoder(mp4, seek_mode="exact").get_frames_at(indices=[5, 0, 5]) @@ -418,3 +418,52 @@ def test_replica_decoders_share_and_release_timeline(tmp_path, variable): finally: collator.close() assert reference() is None + + +def test_nominal_frame_rate_differs_from_presentation_timestamps(): + av = pytest.importorskip("av") + VideoDecoder = require_torchcodec() + output = io.BytesIO() + with av.open(output, "w", format="mp4") as container: + stream = container.add_stream("libx265", rate=60, options={"x265-params": "log-level=error:pools=1"}) + stream.width, stream.height = 64, 48 + stream.pix_fmt = "yuv420p" + stream.codec_context.gop_size = 12 + stream.codec_context.max_b_frames = 0 + for i in range(60): + pixels = np.full((48, 64, 3), i * 4, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(pixels, format="rgb24") + frame.pts, frame.time_base = i, Fraction(1, 60) + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + # Preserve codec headers advertising 60 fps, but remux using capture + # timestamps near 20 fps with variable intervals, as in bag imports. + remuxed = io.BytesIO() + with av.open(io.BytesIO(output.getvalue())) as source: + with av.open(remuxed, "w", format="mp4") as target: + stream = target.add_stream_from_template(source.streams.video[0]) + stream.time_base = Fraction(1, 1000000) + for i, packet in enumerate(p for p in source.demux(video=0) if p.size): + packet.pts = packet.dts = i * 50000 + (i % 2) * 10000 + packet.duration = 50000 + packet.time_base = Fraction(1, 1000000) + packet.stream = stream + target.mux(packet) + payload = remuxed.getvalue() + exact = VideoDecoder(payload, seek_mode="exact") + assert exact.metadata.average_fps_from_header != pytest.approx(20, abs=1) + pts, unused, scale, unused = _timeline(io.BytesIO(payload), len(payload)) + np.testing.assert_allclose(pts / scale, np.arange(60) * 0.05 + np.arange(60) % 2 * 0.01) + source, unused = indexed_source(payload) + decoder = _IndexedTorchCodecDecoder(source, source.video_keyframe_index) + try: + for positions in ([30, 31, 32], [59, 0, 30, 30, 1]): + actual = decoder.get_frames_at(indices=positions) + expected = exact.get_frames_at(indices=positions) + np.testing.assert_array_equal(actual.data.numpy(), expected.data.numpy()) + np.testing.assert_allclose(actual.pts_seconds.numpy(), expected.pts_seconds.numpy()) + assert decoder._indexed is supports_file_like() + finally: + decoder.close()