diff --git a/docs/docs/pypaimon/lerobot.md b/docs/docs/pypaimon/lerobot.md index 518a5d067400..3499a703b76a 100644 --- a/docs/docs/pypaimon/lerobot.md +++ b/docs/docs/pypaimon/lerobot.md @@ -213,8 +213,9 @@ For map-style training, pass an image- or video-backed table group created by ```python from pypaimon.multimodal import PaimonLeRobotDataset +table = conn.get_table("robot_data") dataset = PaimonLeRobotDataset( - conn.get_table("robot_data"), + table, tag_name=tag, ) loader = dataset.to_dataloader( @@ -232,6 +233,10 @@ 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. +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. + Subclass `PaimonDatasetReader` for a custom logical frame layout: ```python diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index 525b16f70892..46322f8fc073 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -1163,6 +1163,15 @@ class CoreOptions: .with_description("Read batch size for any file format if it supports.") ) + READ_VIDEO_MAX_OPEN_DECODERS: ConfigOption[int] = ( + ConfigOptions.key("read.video.max-open-decoders") + .int_type() + .default_value(16) + .with_description( + "Maximum cached video decoders per video feature in a PyPaimon " + "LeRobot DataLoader worker.") + ) + PARQUET_COLUMN_INDEX_ENABLED: ConfigOption[bool] = ( ConfigOptions.key("parquet.filter.columnindex.enabled") .boolean_type() @@ -1927,6 +1936,20 @@ def local_cache_whitelist(self) -> str: def read_batch_size(self, default=None) -> int: 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 + raw = self.options.to_map().get(option.key()) + if isinstance(raw, (bool, float)): + raise ValueError("%s must be a positive integer." % option.key()) + try: + value = self.options.get(option) + except (TypeError, ValueError) as error: + raise ValueError( + "%s must be a positive integer." % option.key()) from error + if value <= 0: + raise ValueError("%s must be a positive integer." % option.key()) + return value + def parquet_column_index_enabled(self) -> bool: return self.options.get(CoreOptions.PARQUET_COLUMN_INDEX_ENABLED) diff --git a/paimon-python/pypaimon/multimodal/lerobot/dataset.py b/paimon-python/pypaimon/multimodal/lerobot/dataset.py index 665ceb147863..38793e26da4e 100644 --- a/paimon-python/pypaimon/multimodal/lerobot/dataset.py +++ b/paimon-python/pypaimon/multimodal/lerobot/dataset.py @@ -34,6 +34,7 @@ import pyarrow as pa +from pypaimon.common.options.core_options import CoreOptions from pypaimon.multimodal.lerobot.metadata import ( _companion_table_identifiers, _restore_pandas_metadata, @@ -267,6 +268,12 @@ def _set_frame_rows( if self._video_keys and self._file_io is None: raise ValueError( "A video-backed PaimonDatasetReader must expose file_io.") + table_options = getattr(self._read_table, "options", None) + cache_size = ( + table_options.read_video_max_open_decoders() + if isinstance(table_options, CoreOptions) + else CoreOptions.READ_VIDEO_MAX_OPEN_DECODERS.default_value() + ) self._video_collators = [ VideoFrameCollator( access, @@ -277,6 +284,7 @@ def _set_frame_rows( output_column=key, collate_fn=_identity, range_parallelism=self.blob_parallelism, + max_open_videos=cache_size, ) for key in self._video_keys ] diff --git a/paimon-python/pypaimon/multimodal/table.py b/paimon-python/pypaimon/multimodal/table.py index 00f1c03c4e60..859a9036dd80 100644 --- a/paimon-python/pypaimon/multimodal/table.py +++ b/paimon-python/pypaimon/multimodal/table.py @@ -102,6 +102,12 @@ def __init__(self, catalog, identifier: str, raw_table): self.raw_table = raw_table self.table = raw_table + def copy(self, options: dict) -> 'MultimodalTable': + """Return a table view with dynamic Paimon options.""" + raw_table = self.raw_table.copy(options) + return MultimodalTable( + self.catalog, raw_table.identifier.get_full_name(), raw_table) + def add(self, data): arrow_table = _to_arrow_table(data, _target_schema(self.raw_table)) write_builder = self.raw_table.new_batch_write_builder() diff --git a/paimon-python/pypaimon/multimodal/video.py b/paimon-python/pypaimon/multimodal/video.py index 05e70619d1aa..90489a0b0a78 100644 --- a/paimon-python/pypaimon/multimodal/video.py +++ b/paimon-python/pypaimon/multimodal/video.py @@ -150,7 +150,7 @@ def __init__( decoder_factory, decode_fn=None, output_column="frame", - max_open_videos=8, + max_open_videos=16, collate_fn=None, decode_batch_fn=None, range_parallelism=4): diff --git a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py index f39d7e1add02..a9eb32ef00fe 100644 --- a/paimon-python/pypaimon/tests/multimodal_lerobot_test.py +++ b/paimon-python/pypaimon/tests/multimodal_lerobot_test.py @@ -42,6 +42,7 @@ import pypaimon.multimodal as pmm from pypaimon.common.identifier import Identifier from pypaimon.common.options import Options +from pypaimon.common.options.core_options import CoreOptions from pypaimon.multimodal.source_utils import _SourceFileIO from pypaimon.multimodal.connection import MultimodalConnection from pypaimon.multimodal.lerobot import load_from_lerobot @@ -56,6 +57,7 @@ _image_tensor, _index_names, _open_video_decoder, + _reader_metadata, _selected_episodes, _stack_visual_windows, _torch_row, @@ -669,6 +671,69 @@ def test_indexed_default_backend_falls_back_on_ffmpeg_error(self): self.assertIs(expected, decoder.get_frames_at(indices=[3])) stream.seek.assert_called_once_with(0) + def test_table_option_limits_video_decoders(self): + from pypaimon.filesystem.local_file_io import LocalFileIO + + info = { + "codebase_version": "v3.0", "fps": 10, + "total_frames": 1, "total_episodes": 1, "total_tasks": 1, + "features": { + "index": {"dtype": "int64", "shape": [1]}, + "episode_index": {"dtype": "int64", "shape": [1]}, + "frame_index": {"dtype": "int64", "shape": [1]}, + "timestamp": {"dtype": "float32", "shape": [1]}, + "task_index": {"dtype": "int64", "shape": [1]}, + "camera": {"dtype": "video", "shape": [2, 2, 3], + "names": ["height", "width", "channels"]}, + "camera_b": {"dtype": "video", "shape": [2, 2, 3], + "names": ["height", "width", "channels"]}, + }, + } + metadata = _reader_metadata({ + "repo_id": "test/video-cache-option", "info": info, + "episodes": [{"episode_index": 0, "dataset_from_index": 0, + "dataset_to_index": 1, "length": 1, + "tasks": ["pick"]}], + "tasks": ["pick"], + }) + module = "pypaimon.multimodal.lerobot.dataset." + option = "read.video.max-open-decoders" + + for configured, expected in ((None, 16), ("2", 2)): + with self.subTest(configured=configured): + options = {} if configured is None else {option: configured} + frames = SimpleNamespace( + options=CoreOptions(Options(options))) + rows = SimpleNamespace( + _table=frames, num_rows=1, file_io=LocalFileIO(), + close=lambda: None) + with patch(module + "_load_dataset", return_value=(frames, metadata)), \ + patch(module + "_target_schema", + return_value=_schema_from_info(info)), \ + patch(module + "_PaimonTableFrameReader", + return_value=rows): + dataset = pmm.PaimonLeRobotDataset(Mock()) + try: + self.assertEqual([expected, expected], [ + c.max_open_videos for c in dataset.reader._video_collators]) + finally: + dataset.reader.close() + + for configured in ("0", "bad", True, 1.5): + with self.subTest(configured=configured): + frames = SimpleNamespace( + options=CoreOptions(Options({option: configured}))) + rows = SimpleNamespace( + _table=frames, num_rows=1, file_io=LocalFileIO(), + close=lambda: None) + with patch(module + "_load_dataset", return_value=(frames, metadata)), \ + patch(module + "_target_schema", + return_value=_schema_from_info(info)), \ + patch(module + "_PaimonTableFrameReader", + return_value=rows): + with self.assertRaisesRegex(ValueError, option): + pmm.PaimonLeRobotDataset(Mock()) + def test_video_batches_include_delta_frames_and_preserve_backends(self): try: import torch @@ -3475,10 +3540,18 @@ def test_table_dataset_pickle_preserves_episode_metadata_and_reads(self): import torch self.connection.load_from_lerobot("worker_pickle", self.image_source) - table = self.connection.get_table("worker_pickle") - dataset = pmm.PaimonLeRobotDataset(table, return_uint8=True) + snapshots = self.connection.create_lerobot_tag( + "worker_pickle", "training") + table = self.connection.get_table("worker_pickle").copy({ + "read.video.max-open-decoders": "2"}) + dataset = pmm.PaimonLeRobotDataset( + table, tag_name="training", return_uint8=True) restored = pickle.loads(pickle.dumps(dataset)) + for reader in (dataset.reader, restored.reader): + self.assertEqual(snapshots["frames"], reader._snapshot_id) + self.assertEqual( + 2, reader._read_table.options.read_video_max_open_decoders()) self.assertEqual(dataset.meta.episodes[:], restored.meta.episodes[:]) self.assertEqual( dataset.meta.episodes._fingerprint, diff --git a/paimon-python/pypaimon/tests/multimodal_table_test.py b/paimon-python/pypaimon/tests/multimodal_table_test.py index 00bd51e7d214..e2f665e8b1a6 100644 --- a/paimon-python/pypaimon/tests/multimodal_table_test.py +++ b/paimon-python/pypaimon/tests/multimodal_table_test.py @@ -98,6 +98,29 @@ def test_connect_rejects_positional_warehouse(self): with self.assertRaises(TypeError): pmm.connect(self.warehouse) + def test_copy_applies_dynamic_table_options(self): + table = self.conn.create_table( + "frames", schema=_schema({"id": pa.int32()}), + options=_PARQUET_OPTIONS) + copied = table.copy({"read.video.max-open-decoders": "8"}) + + self.assertEqual(table.identifier, copied.identifier) + self.assertIs(table.catalog, copied.catalog) + self.assertNotIn( + "read.video.max-open-decoders", + table.raw_table.table_schema.options) + self.assertEqual( + "8", copied.raw_table.table_schema.options[ + "read.video.max-open-decoders"]) + + table.add([{"id": 1}]) + snapshot_id = table.raw_table.snapshot_manager().get_latest_snapshot().id + table.raw_table.create_tag("training", snapshot_id=snapshot_id) + tagged = copied.raw_table.copy({"scan.tag-name": "training"}) + self.assertEqual( + "8", tagged.table_schema.options[ + "read.video.max-open-decoders"]) + def test_create_table_defaults_data_evolution_options(self): table = self.conn.create_table( "docs",