From 6b54aa031c6071e0b9ad5c0b588cd8e5256ae16e Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 26 Aug 2026 14:51:34 +0800 Subject: [PATCH 1/3] feat(python): add lazy contiguous window dataset Add a generic snapshot-pinned window Dataset that builds a lightweight row-id index, reads projected payloads on demand, and handles continuity, tail policies, transforms, and multi-worker DataLoader use. Co-Authored-By: Codex Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex AI-Contributed/Feature: 2/2 AI-Contributed/UT: 0/0 --- docs/docs/pypaimon/multimodal-api.mdx | 58 +++ docs/docs/pypaimon/pytorch.md | 39 ++ paimon-python/pypaimon/multimodal/query.py | 40 ++ .../pypaimon/multimodal/window_dataset.py | 360 ++++++++++++++++++ .../tests/contiguous_window_dataset_test.py | 314 +++++++++++++++ 5 files changed, 811 insertions(+) create mode 100644 paimon-python/pypaimon/multimodal/window_dataset.py create mode 100644 paimon-python/pypaimon/tests/contiguous_window_dataset_test.py diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e1074ee5ce8c..f173afc89fc8 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -720,6 +720,64 @@ Notes: few large reads); scattered point reads coalesce less. - Blob reads are available only on `scan()`, not on the `search()` queries. +### Contiguous windows for PyTorch + +Install the `torch` extra, then use `to_contiguous_window_dataset` to expose +map-style windows without loading the selected rows or BLOB payloads into Python +memory up front. The Dataset builds a compact index from the group column, order +column, and Paimon row IDs. Each `__getitem__` call fetches only that window from +the snapshot recorded in `dataset.snapshot_id`. + +```shell +pip install pypaimon[torch] +``` + +```python +import torch + + +def float32_window(values): + return torch.tensor(values, dtype=torch.float32) + + +windows = ( + frames.scan() + .where("split = 'train'") + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "action"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + column_transforms={ + "state": float32_window, + "action": float32_window, + }, + ) +) + +sample = windows[0] +assert sample["action"].shape == (16, action_size) +assert sample["is_pad"].shape == (16,) +``` + +The group and order keys in a sample identify the window anchor. Every projected +column contains the whole window. With `tail="drop"`, only full windows are +exposed. With `tail="pad"`, every real row is an anchor; missing suffix values +repeat the last real value by default and `is_pad` is `True` exactly at those +positions. With `tail="error"`, construction fails if any scheduled anchor is +incomplete. Use `pad_values` to override the repeated value for individual +columns. Anchors advance by `stride`, which defaults to one row. + +`column_transforms` receive one padded Python list per projected column. This is +where applications define tensor dtype and shape or decode BLOB bytes. The +optional `adapter` receives the resulting sample mapping and can rename or +combine fields for a model-specific batch contract. The core Dataset does not +know model field names, image formats, or normalization rules. Top-level +functions and callable classes are recommended for transforms and adapters so +the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader` +instances. + ### Distributed BLOB processing with Ray For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index ac9f6b8d64eb..2e6e3cfc9994 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -106,7 +106,46 @@ embedded frame ordinals keep frame mapping out of the normal data file. Use physical video ranges and cache decoder sessions per worker. See [Multimodal API: Video Frame Storage](multimodal-api#video-frame-storage) for the write path and a complete decoder example. +## Contiguous Windows +Use a map-style `ContiguousWindowDataset` when training samples are fixed-size +windows which must not cross a sequence boundary. The dataset builds an index +from only the group column, order column, and Paimon row IDs. Projected values, +including BLOB payloads, are read from the pinned snapshot when a sample is +requested; they are not retained in the index. + +```python +from torch.utils.data import DataLoader + +dataset = ( + frames.scan() + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "image"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + ) +) + +loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True) +``` + +Each item contains the group and order keys, one list for each requested +column, and a boolean `is_pad` tensor where `True` marks padding. Padding +repeats the final real value by default; `pad_values` can override individual +columns. Use `column_transforms` to convert column lists to tensors and +`adapter` to produce a model-specific sample mapping. Keep these callbacks +picklable when using multiple DataLoader workers. + +Scheduled anchors start at row zero and advance by `stride` (default `1`). +`tail="drop"` omits incomplete windows, `tail="pad"` includes and pads them, +and `tail="error"` rejects a sequence with any scheduled incomplete window. +Rows are sorted by `order_key` inside each `group_key` value. Order values must +be integers which increase by exactly one; duplicates and missing steps are +rejected, and windows never cross groups. The resolved Paimon +snapshot is pinned for the lifetime of the dataset, so later commits cannot +change its index or sample contents. ## File Format Metadata Cache Reusable PyArrow Dataset metadata is cached across reads. Configure its estimated diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 9046498f7cf8..4b614de11c49 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -151,6 +151,46 @@ def to_torch( max_buffer_input_splits=max_buffer_input_splits, ) + def to_contiguous_window_dataset( + self, + *, + window_size, + columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + """Build a snapshot-pinned, map-style Dataset of contiguous rows. + + The Dataset indexes only ``group_key``, ``order_key``, and Paimon row + IDs, then reads projected values on demand. It sorts rows within each + group and never creates a window across groups. See + :class:`pypaimon.multimodal.window_dataset.ContiguousWindowDataset` + for tail, padding, mask, transform, and adapter semantics. + """ + if self._result_factory is not None: + raise TypeError( + "to_contiguous_window_dataset is only supported on scan(), " + "not search queries.") + from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + return ContiguousWindowDataset( + self, + window_size=window_size, + columns=columns, + group_key=group_key, + order_key=order_key, + stride=stride, + tail=tail, + column_transforms=column_transforms, + pad_values=pad_values, + adapter=adapter, + blob_parallelism=blob_parallelism, + ) + def to_ray( self, *, diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py new file mode 100644 index 000000000000..6517d9e1d1e4 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -0,0 +1,360 @@ +# 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. + +"""Snapshot-pinned PyTorch Dataset for contiguous Paimon row windows.""" + +import copy +import operator +from collections import defaultdict +from numbers import Integral + +import torch +from torch.utils.data import Dataset + +from pypaimon.common.options.core_options import CoreOptions +from pypaimon.multimodal.query import ScanQuery +from pypaimon.schema.data_types import is_blob_type +from pypaimon.snapshot.time_travel_util import SCAN_KEYS +from pypaimon.table.special_fields import SpecialFields + + +class ContiguousWindowDataset(Dataset): + """Map-style Dataset which reads fixed row windows on demand. + + The in-memory index contains only group values, order values, and Paimon + row IDs. Each ``__getitem__`` reads the projected rows from the snapshot + resolved while the index was built. ``tail`` controls scheduled anchors + whose remaining rows are shorter than ``window_size``: + + * ``drop`` omits them; + * ``pad`` repeats final values and marks repeats in ``is_pad``; + * ``error`` rejects the dataset. + + ``column_transforms`` convert individual padded column lists and + ``adapter`` can adapt the complete mapping to a model-specific contract. + """ + + _TAIL_POLICIES = ("drop", "pad", "error") + + def __init__( + self, + query, + *, + window_size, + columns=None, + group_key="episode_id", + order_key="step_idx", + stride=1, + tail="drop", + column_transforms=None, + pad_values=None, + adapter=None, + blob_parallelism=64): + if getattr(query, "_result_factory", None) is not None: + raise TypeError( + "ContiguousWindowDataset is only supported on scan(), " + "not search queries.") + self.window_size = _positive_int(window_size, "window_size") + self.stride = _positive_int(stride, "stride") + if tail not in self._TAIL_POLICIES: + raise ValueError( + "tail must be one of %s; got %r." + % (self._TAIL_POLICIES, tail)) + self.tail = tail + self.group_key = _column(query, group_key, "group_key") + self.order_key = _column(query, order_key, "order_key") + if self.group_key == self.order_key: + raise ValueError("group_key and order_key must name different columns.") + if "is_pad" in (self.group_key, self.order_key): + raise ValueError("group_key and order_key must not be is_pad.") + self.columns = _columns( + query, columns, self.group_key, self.order_key) + self.column_transforms = _column_transforms( + column_transforms, self.columns) + self.pad_values = _pad_values(pad_values, self.columns) + if adapter is not None and not callable(adapter): + raise TypeError("adapter must be callable or None.") + self.adapter = adapter + self.blob_parallelism = _positive_int( + blob_parallelism, "blob_parallelism") + + if not query._table.options.row_tracking_enabled(): + raise ValueError( + "ContiguousWindowDataset requires row-tracking.enabled=true.") + + self._blob_columns = [ + field.name for field in query._table.fields + if field.name in self.columns and is_blob_type(field.type) + ] + index, snapshot_id = _read_window_index( + query, self.group_key, self.order_key) + self.snapshot_id = snapshot_id + self._table = _pin_table(query._table, snapshot_id) + self._groups, self._anchors = self._build_index(index) + + @classmethod + def from_query(cls, query, **kwargs): + """Build a contiguous-window Dataset from a ``ScanQuery``.""" + return cls(query, **kwargs) + + def __len__(self): + return len(self._anchors) + + def __getitem__(self, index): + index = operator.index(index) + if index < 0: + index += len(self._anchors) + if index < 0 or index >= len(self._anchors): + raise IndexError("window index out of range") + + group_index, start, valid_count = self._anchors[index] + group_key, order_values, row_ids = self._groups[group_index] + selected_row_ids = row_ids[start:start + valid_count] + rows = self._read_rows(selected_row_ids) + padding_count = self.window_size - valid_count + padding_mask = torch.zeros(self.window_size, dtype=torch.bool) + if padding_count: + padding_mask[valid_count:] = True + sample = { + self.group_key: group_key, + self.order_key: order_values[start], + "is_pad": padding_mask, + } + for name in self.columns: + values = [row[name] for row in rows] + if padding_count: + pad_value = self.pad_values.get(name, values[-1]) + values.extend( + copy.deepcopy(pad_value) for _ in range(padding_count)) + transform = self.column_transforms.get(name) + sample[name] = transform(values) if transform is not None else values + if self.adapter is not None: + return self.adapter(sample) + return sample + + def _build_index(self, index): + group_values = index.column(self.group_key).to_pylist() + order_values = index.column(self.order_key).to_pylist() + row_ids = index.column(SpecialFields.ROW_ID.name).to_pylist() + grouped = defaultdict(list) + for group_key, order_value, row_id in zip( + group_values, order_values, row_ids): + if group_key is None: + raise ValueError("%s must not contain null values." % self.group_key) + if order_value is None: + raise ValueError("%s must not contain null values." % self.order_key) + if isinstance(order_value, bool) or not isinstance(order_value, Integral): + raise ValueError( + "%s must contain integer values." % self.order_key) + try: + grouped[group_key].append((int(order_value), int(row_id))) + except TypeError: + raise ValueError( + "%s values must be hashable." % self.group_key) + + groups = [] + anchors = [] + try: + sorted_groups = sorted(grouped.items(), key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values must be mutually orderable." % self.group_key) + for group_key, members in sorted_groups: + try: + members.sort(key=lambda item: item[0]) + except TypeError: + raise ValueError( + "%s values in group %r must be mutually orderable." + % (self.order_key, group_key)) + for previous, current in zip(members, members[1:]): + if previous[0] == current[0]: + raise ValueError( + "Group %s has duplicate order value %r in %s." + % (group_key, current[0], self.order_key)) + if current[0] != previous[0] + 1: + raise ValueError( + "Group %s is not contiguous in %s: %s followed by %s." + % (group_key, self.order_key, + previous[0], current[0])) + + group_index = len(groups) + group_orders = [member[0] for member in members] + group_row_ids = [member[1] for member in members] + groups.append((group_key, group_orders, group_row_ids)) + for start in range(0, len(members), self.stride): + valid_count = min(self.window_size, len(members) - start) + if valid_count < self.window_size: + if self.tail == "drop": + continue + if self.tail == "error": + raise ValueError( + "Group %s has an incomplete window at %s: " + "window_size=%d, available=%d." + % (group_key, group_orders[start], + self.window_size, valid_count)) + anchors.append((group_index, start, valid_count)) + return groups, anchors + + def _read_rows(self, row_ids): + query = ScanQuery(self._table) + predicate_builder = ( + self._table.new_read_builder() + .with_projection( + [field.name for field in self._table.fields] + + [SpecialFields.ROW_ID.name]) + .new_predicate_builder() + ) + query._predicate = predicate_builder.is_in( + SpecialFields.ROW_ID.name, row_ids) + query._projection = list(self.columns) + query._include_row_id = True + + if self._blob_columns: + scalar, blobs = query.read_blobs( + self._blob_columns, parallelism=self.blob_parallelism) + rows = scalar.to_pylist() + for name in self._blob_columns: + values = blobs[name] + if len(values) != len(rows): + raise RuntimeError( + "BLOB column %s is not row-aligned with a window read." + % name) + for row, value in zip(rows, values): + row[name] = value + else: + rows = query.to_arrow().to_pylist() + + by_row_id = {} + row_id_column = SpecialFields.ROW_ID.name + for row in rows: + row_id = int(row[row_id_column]) + del row[row_id_column] + by_row_id[row_id] = row + missing = [row_id for row_id in row_ids if row_id not in by_row_id] + if missing: + raise RuntimeError( + "Pinned snapshot %s did not return indexed row IDs %s." + % (self.snapshot_id, missing)) + return [by_row_id[row_id] for row_id in row_ids] + + +def _read_window_index(query, group_by, order_by): + index_query = copy.copy(query) + index_query._projection = [group_by, order_by] + index_query._include_row_id = True + read_builder = index_query._configured_read_builder() + plan = read_builder.new_scan().plan() + index = read_builder.new_read().to_arrow(plan.splits()) + if index.num_rows and plan.snapshot_id is None: + raise RuntimeError("Cannot pin the snapshot used to build the window index.") + return index, plan.snapshot_id + + +def _pin_table(table, snapshot_id): + if snapshot_id is None: + return table + options = { + key: None for key in SCAN_KEYS + if table.options.options.contains_key(key) + } + options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) + return table.copy(options) + + +def _columns(query, columns, group_key, order_key): + available = {field.name for field in query._table.fields} + if columns is None: + if query._projection is None: + columns = [field.name for field in query._table.fields] + else: + columns = list(query._projection) + columns = [name for name in columns + if name not in (group_key, order_key)] + elif isinstance(columns, str): + columns = [columns] + else: + try: + columns = list(columns) + except TypeError: + raise TypeError( + "columns must be a non-empty sequence of column names.") + if not columns: + raise ValueError("columns must contain at least one value column.") + if any(not isinstance(name, str) or not name for name in columns): + raise TypeError("columns must contain only non-empty column names.") + if len(set(columns)) != len(columns): + raise ValueError("columns must not contain duplicates.") + invalid = [name for name in columns if name not in available] + if invalid: + raise ValueError("columns do not exist: %s." % invalid) + reserved = [name for name in columns + if name in (group_key, order_key, "is_pad")] + if reserved: + raise ValueError( + "columns must not include group_key, order_key, or is_pad: %s." + % reserved) + return columns + + +def _column_transforms(value, columns): + transforms = _mapping(value, "column_transforms") + _validate_mapping_columns(transforms, columns, "column_transforms") + invalid = [name for name, transform in transforms.items() + if not callable(transform)] + if invalid: + raise TypeError( + "column_transforms values must be callable: %s." % invalid) + return transforms + + +def _pad_values(value, columns): + values = _mapping(value, "pad_values") + _validate_mapping_columns(values, columns, "pad_values") + return values + + +def _mapping(value, name): + if value is None: + return {} + try: + return dict(value) + except (TypeError, ValueError): + raise TypeError("%s must be a mapping or None." % name) + + +def _validate_mapping_columns(value, columns, name): + invalid = [column for column in value if column not in columns] + if invalid: + raise ValueError("%s contains unknown columns: %s." % (name, invalid)) + + +def _column(query, value, name): + if not isinstance(value, str) or not value: + raise TypeError("%s must be a non-empty column name." % name) + available = {field.name for field in query._table.fields} + if value not in available: + raise ValueError("%s column %r does not exist." % (name, value)) + return value + + +def _positive_int(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("%s must be a positive int." % name) + return value + + +__all__ = ["ContiguousWindowDataset"] diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py new file mode 100644 index 000000000000..34ce7e647d1a --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -0,0 +1,314 @@ +# 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 os +import shutil +import tempfile +import unittest +from unittest.mock import patch + +import pyarrow as pa +import torch + +import pypaimon.multimodal as pmm +from pypaimon.multimodal.query import ScanQuery +from pypaimon.multimodal.window_dataset import ContiguousWindowDataset + + +_TABLE_OPTIONS = { + "row-tracking.enabled": "true", + "data-evolution.enabled": "true", + "deletion-vectors.enabled": "true", + "file.format": "parquet", + "vector.file.format": "parquet", +} + + +class _TensorColumnTransform: + + def __call__(self, values): + return torch.tensor(values, dtype=torch.int64) + + +class _WindowAdapter: + + def __call__(self, sample): + return { + "episode": sample["episode"], + "start": sample["step"], + "values": sample["value"], + "padding_mask": sample["is_pad"], + } + + +class ContiguousWindowDatasetTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp(prefix="pypaimon_windows_") + self.conn = pmm.connect(options={ + "warehouse": os.path.join(self.temp_dir, "warehouse"), + }) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + @staticmethod + def _schema(): + return pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + pa.field("payload", pa.large_binary(), nullable=False), + ]) + + @staticmethod + def _row(episode, step): + return { + "episode": episode, + "step": step, + "value": step + (100 if episode == "episode-b" else 0), + "payload": ("%s-%d" % (episode, step)).encode(), + } + + def _table(self, name="frames"): + table = self.conn.create_table( + name, schema=self._schema(), options=_TABLE_OPTIONS) + table.add([ + self._row("episode-b", 2), + self._row("episode-a", 1), + self._row("episode-b", 0), + self._row("episode-a", 0), + self._row("episode-b", 3), + self._row("episode-b", 1), + ]) + return table + + @staticmethod + def _dataset(table, **kwargs): + return ( + table.scan() + .to_contiguous_window_dataset( + window_size=3, + columns=["value", "payload"], + group_key="episode", + order_key="step", + **kwargs, + ) + ) + + def test_sorts_rows_and_never_crosses_episode_boundaries(self): + dataset = self._dataset(self._table()) + + self.assertIsInstance(dataset, torch.utils.data.Dataset) + self.assertEqual(2, len(dataset)) + self.assertIsInstance(dataset.snapshot_id, int) + self.assertNotIn("_episodes", vars(dataset)) + + first = dataset[0] + second = dataset[1] + self.assertEqual("episode-b", first["episode"]) + self.assertEqual(0, first["step"]) + self.assertEqual([100, 101, 102], first["value"]) + self.assertEqual([101, 102, 103], second["value"]) + self.assertFalse(first["is_pad"].any()) + self.assertEqual({"episode-b"}, { + window["episode"] for window in (first, second) + }) + + def test_reads_blob_payloads_only_when_a_window_is_requested(self): + table = self._table() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + dataset = self._dataset(table) + self.assertEqual(0, fetch.call_count) + + sample = dataset[0] + + self.assertEqual(1, fetch.call_count) + self.assertEqual(3, len(fetch.call_args.args[1]["payload"])) + self.assertEqual( + [b"episode-b-0", b"episode-b-1", b"episode-b-2"], + sample["payload"], + ) + + def test_pad_tail_repeats_last_row_and_marks_real_padding(self): + dataset = self._dataset( + self._table(), tail="pad", pad_values={"value": -1}) + + self.assertEqual(6, len(dataset)) + short_tail = dataset[1] + long_tail = dataset[-1] + self.assertEqual("episode-a", short_tail["episode"]) + self.assertEqual([1, -1, -1], short_tail["value"]) + self.assertEqual( + [b"episode-a-1"] * 3, short_tail["payload"]) + self.assertEqual([False, True, True], short_tail["is_pad"].tolist()) + self.assertEqual("episode-b", long_tail["episode"]) + self.assertEqual([103, -1, -1], long_tail["value"]) + self.assertEqual([False, True, True], long_tail["is_pad"].tolist()) + + def test_error_tail_rejects_an_incomplete_scheduled_window(self): + with self.assertRaisesRegex( + ValueError, "episode-a.*incomplete.*window_size=3"): + self._dataset(self._table(), tail="error") + + def test_stride_controls_scheduled_window_anchors(self): + dataset = self._dataset(self._table(), stride=2, tail="pad") + + self.assertEqual( + [("episode-a", 0), ("episode-b", 0), ("episode-b", 2)], + [(dataset[index]["episode"], dataset[index]["step"]) + for index in range(len(dataset))], + ) + self.assertEqual( + [False, False, True], dataset[-1]["is_pad"].tolist()) + + def test_rejects_missing_and_duplicate_order_keys_within_a_group(self): + gapped = self.conn.create_table( + "gapped", schema=self._schema(), options=_TABLE_OPTIONS) + gapped.add([ + self._row("episode-a", 0), + self._row("episode-a", 2), + ]) + + with self.assertRaisesRegex( + ValueError, "episode-a.*not contiguous.*0.*2"): + self._dataset(gapped) + + table = self.conn.create_table( + "duplicates", schema=self._schema(), options=_TABLE_OPTIONS) + table.add([ + self._row("episode-a", 0), + self._row("episode-a", 0), + self._row("episode-a", 1), + ]) + + with self.assertRaisesRegex( + ValueError, "episode-a.*duplicate.*order.*0"): + self._dataset(table) + + def test_pins_snapshot_for_later_on_demand_reads(self): + table = self._table() + dataset = self._dataset(table) + snapshot_id = dataset.snapshot_id + + table.add([self._row("episode-b", 4)]) + + self.assertEqual(snapshot_id, dataset.snapshot_id) + self.assertNotEqual( + snapshot_id, table.raw_table.snapshot_manager().get_latest_snapshot().id) + self.assertEqual(2, len(dataset)) + self.assertEqual([101, 102, 103], dataset[-1]["value"]) + + def test_projection_filter_transform_and_dataloader_workers(self): + table = self._table() + dataset = ( + table.scan() + .where("episode = 'episode-b'") + .select(["value"]) + .to_contiguous_window_dataset( + window_size=2, + group_key="episode", + order_key="step", + column_transforms={"value": _TensorColumnTransform()}, + adapter=_WindowAdapter(), + ) + ) + + loader = torch.utils.data.DataLoader( + dataset, batch_size=2, shuffle=False, num_workers=2) + batches = list(loader) + + self.assertEqual(2, len(batches)) + self.assertEqual(torch.int64, batches[0]["values"].dtype) + self.assertEqual((2, 2), tuple(batches[0]["values"].shape)) + self.assertEqual(torch.bool, batches[0]["padding_mask"].dtype) + self.assertEqual([0, 1, 2], [ + start for batch in batches for start in batch["start"].tolist() + ]) + self.assertEqual( + [[100, 101], [101, 102], [102, 103]], + [values for batch in batches for values in batch["values"].tolist()], + ) + self.assertTrue(all( + episode == "episode-b" + for batch in batches for episode in batch["episode"] + )) + + def test_default_keys_and_public_from_query_entry_point(self): + table = self.conn.create_table( + "default_keys", + schema=pa.schema([ + pa.field("episode_id", pa.string(), nullable=False), + pa.field("step_idx", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), + options=_TABLE_OPTIONS, + ) + table.add([ + {"episode_id": "episode-a", "step_idx": 0, "value": 10}, + {"episode_id": "episode-a", "step_idx": 1, "value": 11}, + ]) + + dataset = ContiguousWindowDataset.from_query( + table.scan().select(["value"]), window_size=2) + + self.assertEqual(1, len(dataset)) + self.assertEqual("episode-a", dataset[0]["episode_id"]) + self.assertEqual(0, dataset[0]["step_idx"]) + self.assertEqual([10, 11], dataset[0]["value"]) + + def test_validates_configuration_and_scan_only_contract(self): + table = self._table() + query = table.scan() + for name, value in ( + ("window_size", 0), + ("stride", 0), + ("tail", "unknown"), + ("group_key", "missing"), + ("order_key", "missing")): + kwargs = { + "window_size": 2, + "columns": ["value"], + "stride": 1, + "tail": "drop", + "group_key": "episode", + "order_key": "step", + } + kwargs[name] = value + with self.subTest(name=name), self.assertRaises((TypeError, ValueError)): + query.to_contiguous_window_dataset(**kwargs) + + reserved_table = self.conn.create_table( + "reserved", schema=pa.schema([ + pa.field("is_pad", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("value", pa.int32(), nullable=False), + ]), options=_TABLE_OPTIONS) + with self.assertRaisesRegex(ValueError, "must not be is_pad"): + reserved_table.scan().to_contiguous_window_dataset( + window_size=2, columns=["value"], + group_key="is_pad", order_key="step") + + with self.assertRaisesRegex(TypeError, "only supported on scan"): + table.search("anything", column="episode").to_contiguous_window_dataset( + window_size=2, columns=["value"], + group_key="episode", order_key="step") + + +if __name__ == "__main__": + unittest.main() From 9c6270951d6ddc693c4b844b699db21d8f3aba82 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 26 Aug 2026 14:52:44 +0800 Subject: [PATCH 2/3] feat(python): add paired HDF5 Paimon ACT benchmark Run both backends through one deterministic ACT harness and pin lazy Paimon windows to the normalization snapshot. Adapt the benchmark to the rebased lazy source contract, coalesce overlapping batch reads, and measure Python allocations outside wall-clock timing. Co-Authored-By: Codex Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex AI-Contributed/Feature: 1418/1418 AI-Contributed/UT: 363/363 --- docs/docs/pypaimon/robomind-act-benchmark.md | 76 ++ .../pypaimon/benchmark/act_harness.py | 470 ++++++++++ .../pypaimon/benchmark/paired_act.py | 835 ++++++++++++++++++ .../pypaimon/multimodal/window_dataset.py | 33 +- .../tests/contiguous_window_dataset_test.py | 20 + .../tests/paired_act_benchmark_test.py | 343 +++++++ paimon-python/setup.py | 4 + 7 files changed, 1777 insertions(+), 4 deletions(-) create mode 100644 docs/docs/pypaimon/robomind-act-benchmark.md create mode 100644 paimon-python/pypaimon/benchmark/act_harness.py create mode 100644 paimon-python/pypaimon/benchmark/paired_act.py create mode 100644 paimon-python/pypaimon/tests/paired_act_benchmark_test.py diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md new file mode 100644 index 000000000000..0884553b9307 --- /dev/null +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -0,0 +1,76 @@ +--- +title: "RoboMIND Paired ACT Benchmark" +sidebar_position: 8 +--- + + + +# RoboMIND Paired ACT Benchmark + +The paired benchmark compares original RoboMIND AgileX HDF5 with an already +ingested and canonical-action-backfilled Paimon warehouse. It does not include +ingestion or backfill time. Install the ACT and HDF5 extras, run the +[RoboMIND AgileX pipeline](robomind-agilex), and then execute: + +```shell +pip install 'pypaimon[act,hdf5]' +python -m pypaimon.benchmark.paired_act \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --warehouse /data/warehouse \ + --report /data/results/paired-act.json +``` + +One immutable configuration controls both paths. The runner computes train-only +normalization once, verifies its canonical action values against the requested +version in `feature_stats_agilex`, and passes the same object to both adapters. +A seeded window plan fixes every warmup, loader, training, and validation +anchor. Before training, the runner requires exact `torch.equal` parity for +sample identity, state, action, image, and padding tensors. + +The Paimon adapter uses `ContiguousWindowDataset`, not an ACT-specific table +reader. Dataset construction indexes only episode, frame, and row IDs. Window +payloads remain lazy until `__getitem__`, and all train and validation reads are +pinned to the exact frames snapshot recorded by the normalization statistics. +PyTorch batch access coalesces overlapping row IDs into one payload read. The +adapter maps each generic window to the same tensor contract as the HDF5 adapter +without materializing episodes in memory. + +Each backend then uses the same CPU LeRobot ACT policy, initial seed, AdamW +optimizer, batch size, window sequence, and optimizer step count. At least +three rounds run in alternating order (`HDF5 → Paimon`, then +`Paimon → HDF5`) to expose ordering effects. The benchmark does not drop the OS +page cache and records `cache_control=uncontrolled`. + +The JSON report contains: + +- input manifest, table snapshot, normalization, configuration, and window + sequence digests; +- exact tensor, train-loss, and validation-loss parity gates; +- first-batch latency, DataLoader samples per second, fixed-step time, and a + separate dataset-build-plus-first-batch Python allocation replay for every + run; +- per-backend median, minimum, and maximum across rounds; +- explicit unverified scope, including native-memory completeness, GPU, + multi-worker loading, distributed training, recovery, and policy quality. + +Python peak allocation uses `tracemalloc` after wall-clock measurement so its +overhead does not distort throughput. The replay covers dataset construction +and one first batch; it does not include every native Arrow or Torch allocation. +Treat it as a reproducible engineering diagnostic, not total process RSS. diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act_harness.py new file mode 100644 index 000000000000..15e5a7d3364c --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act_harness.py @@ -0,0 +1,470 @@ +# 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. + +"""Shared deterministic ACT model, trainer, and window plan for benchmarks.""" + +import gc +import hashlib +import json +import math +import random +import time +import tracemalloc +from dataclasses import asdict, dataclass +from io import BytesIO + +import numpy as np +import torch +import torch.nn.functional as functional +from PIL import Image +from torch.utils.data import DataLoader, Dataset + + +CAMERA_KEYS = ( + "observation.images.front", + "observation.images.left_wrist", + "observation.images.right_wrist", +) + + +@dataclass(frozen=True) +class BenchmarkConfig: + """One immutable ACT and measurement configuration for both backends.""" + + seed: int = 20260825 + action_horizon: int = 32 + batch_size: int = 2 + optimizer_steps: int = 2 + image_height: int = 64 + image_width: int = 80 + learning_rate: float = 1e-4 + weight_decay: float = 1e-4 + warmup_batches: int = 1 + loader_batches: int = 4 + rounds: int = 3 + + def __post_init__(self): + positive_ints = ( + "action_horizon", + "batch_size", + "optimizer_steps", + "image_height", + "image_width", + "warmup_batches", + "loader_batches", + ) + for name in positive_ints: + value = getattr(self, name) + if ( + isinstance(value, bool) + or not isinstance(value, int) + or value <= 0): + raise ValueError("%s must be a positive int." % name) + if isinstance(self.seed, bool) or not isinstance(self.seed, int): + raise ValueError("seed must be an int.") + if isinstance(self.rounds, bool) or not isinstance(self.rounds, int): + raise ValueError("rounds must be an int.") + if self.rounds < 3: + raise ValueError("rounds must be at least 3.") + if self.learning_rate <= 0: + raise ValueError("learning_rate must be positive.") + if self.weight_decay < 0: + raise ValueError("weight_decay must not be negative.") + + def to_dict(self): + return asdict(self) + + +@dataclass(frozen=True) +class WindowPlan: + """Explicit window indices consumed identically by both backends.""" + + seed: int + loader_indices: tuple + train_indices: tuple + validation_indices: tuple + + @property + def sha256(self): + payload = json.dumps( + self.to_dict(), sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def to_dict(self): + return { + "seed": self.seed, + "loader_indices": list(self.loader_indices), + "train_indices": list(self.train_indices), + "validation_indices": list(self.validation_indices), + } + + +def build_window_plan(train_window_count, validation_window_count, config): + """Build stable indices for training, validation, and loader timing.""" + train_window_count = _positive_int( + train_window_count, "train_window_count") + validation_window_count = _positive_int( + validation_window_count, "validation_window_count") + loader_count = ( + config.warmup_batches + config.loader_batches) * config.batch_size + train_count = config.optimizer_steps * config.batch_size + return WindowPlan( + seed=config.seed, + loader_indices=tuple(_repeat_permutations( + train_window_count, loader_count, config.seed + 1)), + train_indices=tuple(_repeat_permutations( + train_window_count, train_count, config.seed + 2)), + validation_indices=tuple(_repeat_permutations( + validation_window_count, config.batch_size, config.seed + 3)), + ) + + +def decode_rgb_image(payload): + """Decode one JPEG/PNG payload identically for HDF5 and Paimon.""" + try: + return np.asarray(Image.open(BytesIO(payload)).convert("RGB")) + except Exception as error: + raise ValueError("Cannot decode ACT RGB image bytes.") from error + + +def validate_act_batch(batch, config): + """Validate the exact tensor contract passed to the shared ACT policy.""" + required = { + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + } + if set(batch) != required: + raise ValueError( + "ACT batch fields differ: expected %s, got %s." + % (sorted(required), sorted(batch))) + batch_size = len(batch["sample_id"]) + expected = { + "qpos": ((batch_size, 14), torch.float32), + "action": ((batch_size, config.action_horizon, 14), torch.float32), + "images": ( + (batch_size, len(CAMERA_KEYS), 3) + + tuple(batch["images"].shape[-2:]), + torch.float32, + ), + "is_pad": ((batch_size, config.action_horizon), torch.bool), + "step_idx": ((batch_size,), torch.int64), + } + for name, (shape, dtype) in expected.items(): + value = batch[name] + if not isinstance(value, torch.Tensor): + raise ValueError("%s must be a torch.Tensor." % name) + if tuple(value.shape) != shape: + raise ValueError( + "%s has shape %s; expected %s." + % (name, tuple(value.shape), shape)) + if value.dtype != dtype: + raise ValueError( + "%s has dtype %s; expected %s." % (name, value.dtype, dtype)) + for name in ("qpos", "action", "images"): + if not torch.isfinite(batch[name]).all(): + raise ValueError("%s contains NaN or Inf." % name) + if torch.any(batch["images"] < 0) or torch.any(batch["images"] > 1): + raise ValueError("images must be normalized to [0, 1].") + if batch["is_pad"].any(): + raise ValueError("M0 ACT windows must be complete and unpadded.") + for sample_id, episode_id, step_idx in zip( + batch["sample_id"], batch["episode_id"], + batch["step_idx"].tolist()): + if sample_id != "%s#%s" % (episode_id, step_idx): + raise ValueError( + "sample_id is not aligned with episode_id and step_idx.") + + +def build_lerobot_batch(batch, config): + """Map the common window contract to LeRobot ACTPolicy feature names.""" + validate_act_batch(batch, config) + images = batch["images"] + target_size = (config.image_height, config.image_width) + if tuple(images.shape[-2:]) != target_size: + flat = images.flatten(0, 1) + flat = functional.interpolate( + flat, size=target_size, mode="bilinear", align_corners=False) + images = flat.reshape(images.shape[:3] + target_size) + result = { + "observation.state": batch["qpos"], + "action": batch["action"], + "action_is_pad": batch["is_pad"], + } + for index, name in enumerate(CAMERA_KEYS): + result[name] = images[:, index] + return result + + +def build_act_policy(config): + """Build the one reduced CPU LeRobot ACT configuration used by M0.""" + try: + import importlib.metadata + from lerobot.configs.types import FeatureType, PolicyFeature + from lerobot.policies.act.configuration_act import ACTConfig + from lerobot.policies.act.modeling_act import ACTPolicy + except ImportError as error: + raise ImportError( + "Paired ACT benchmark requires: " + "pip install -e '.[act]'.") from error + + inputs = { + "observation.state": PolicyFeature(FeatureType.STATE, (14,)), + } + inputs.update({ + name: PolicyFeature( + FeatureType.VISUAL, + (3, config.image_height, config.image_width), + ) + for name in CAMERA_KEYS + }) + act_config = ACTConfig( + input_features=inputs, + output_features={ + "action": PolicyFeature(FeatureType.ACTION, (14,)), + }, + device="cpu", + chunk_size=config.action_horizon, + n_action_steps=config.action_horizon, + vision_backbone="resnet18", + pretrained_backbone_weights=None, + dim_model=64, + n_heads=4, + dim_feedforward=256, + n_encoder_layers=1, + n_decoder_layers=1, + use_vae=True, + latent_dim=16, + n_vae_encoder_layers=1, + kl_weight=10.0, + ) + policy = ACTPolicy(act_config) + return policy, { + "implementation": "lerobot.ACTPolicy", + "lerobot_version": importlib.metadata.version("lerobot"), + "vision_backbone": act_config.vision_backbone, + "pretrained_backbone_weights": act_config.pretrained_backbone_weights, + "chunk_size": act_config.chunk_size, + "dim_model": act_config.dim_model, + "n_heads": act_config.n_heads, + "n_encoder_layers": act_config.n_encoder_layers, + "n_decoder_layers": act_config.n_decoder_layers, + "n_vae_encoder_layers": act_config.n_vae_encoder_layers, + "latent_dim": act_config.latent_dim, + "kl_weight": act_config.kl_weight, + "parameter_count": sum( + parameter.numel() for parameter in policy.parameters()), + "trainable_parameter_count": sum( + parameter.numel() + for parameter in policy.parameters() if parameter.requires_grad), + } + + +def run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sample_sequence_sha256, + policy_factory=None): + """Measure a backend with the shared plan, model, and trainer.""" + _seed_everything(config.seed) + policy_factory = policy_factory or build_act_policy + started = time.monotonic() + dataset_started = time.monotonic() + train_dataset, validation_dataset = dataset_factory() + dataset_build_s = time.monotonic() - dataset_started + + loader_sequence = _SequenceDataset(train_dataset, plan.loader_indices) + loader = DataLoader( + loader_sequence, + batch_size=config.batch_size, + shuffle=False, + num_workers=0, + ) + iterator = iter(loader) + first_batch_started = time.monotonic() + first_batch = next(iterator) + first_batch_s = time.monotonic() - first_batch_started + validate_act_batch(first_batch, config) + for _ in range(config.warmup_batches - 1): + validate_act_batch(next(iterator), config) + + loader_started = time.monotonic() + loader_sample_count = 0 + for _ in range(config.loader_batches): + batch = next(iterator) + validate_act_batch(batch, config) + loader_sample_count += len(batch["sample_id"]) + loader_seconds = time.monotonic() - loader_started + + _seed_everything(config.seed) + policy, model = policy_factory(config) + parameters = ( + policy.get_optim_params() + if hasattr(policy, "get_optim_params") else policy.parameters()) + optimizer = torch.optim.AdamW( + parameters, + lr=config.learning_rate, + weight_decay=config.weight_decay, + ) + policy.train() + training_loader = DataLoader( + _SequenceDataset(train_dataset, plan.train_indices), + batch_size=config.batch_size, + shuffle=False, + num_workers=0, + ) + train_started = time.monotonic() + losses = [] + for step, batch in enumerate(training_loader, 1): + step_started = time.monotonic() + model_batch = build_lerobot_batch(batch, config) + optimizer.zero_grad(set_to_none=True) + loss, components = policy(model_batch) + if loss.ndim != 0 or not torch.isfinite(loss): + raise FloatingPointError( + "ACT produced a non-finite scalar loss at step %d." % step) + loss.backward() + optimizer.step() + losses.append({ + "step": step, + "total": float(loss.detach()), + "components": { + name: _finite_float(value, name) + for name, value in components.items() + }, + "step_time_s": time.monotonic() - step_started, + }) + fixed_steps_s = time.monotonic() - train_started + if len(losses) != config.optimizer_steps: + raise AssertionError( + "Expected %d optimizer steps, got %d." + % (config.optimizer_steps, len(losses))) + + # ACTPolicy only constructs the VAE posterior needed by its supervised + # loss while the module is in training mode. Keep that mode for validation + # but disable gradients and parameter updates below. + policy.train() + _seed_everything(config.seed + 4) + validation_batch = next(iter(DataLoader( + _SequenceDataset(validation_dataset, plan.validation_indices), + batch_size=config.batch_size, + shuffle=False, + num_workers=0, + ))) + with torch.no_grad(): + validation_loss, _ = policy(build_lerobot_batch( + validation_batch, config)) + validation_value = _finite_float(validation_loss, "validation_loss") + wall_time_s = time.monotonic() - started + python_peak = _measure_python_peak(dataset_factory, plan, config) + + return { + "round": round_number, + "backend": backend, + "sample_sequence_sha256": sample_sequence_sha256, + "model": model, + "optimizer": { + "name": "AdamW", + "learning_rate": config.learning_rate, + "weight_decay": config.weight_decay, + }, + "warmup_batches": config.warmup_batches, + "first_batch_s": first_batch_s, + "dataset_build_s": dataset_build_s, + "dataloader_samples": loader_sample_count, + "dataloader_s": loader_seconds, + "dataloader_samples_per_s": loader_sample_count / loader_seconds, + "fixed_steps_s": fixed_steps_s, + "train_loss": [item["total"] for item in losses], + "train_trace": losses, + "validation_loss": validation_value, + "python_peak_allocated_bytes": python_peak, + "peak_memory_measurement": ( + "python-tracemalloc-separate-dataset-first-batch"), + "wall_time_s": wall_time_s, + } + + +def _measure_python_peak(dataset_factory, plan, config): + gc.collect() + tracemalloc.start() + try: + train_dataset, _ = dataset_factory() + indices = plan.loader_indices[:config.batch_size] + next(iter(DataLoader( + _SequenceDataset(train_dataset, indices), + batch_size=config.batch_size, + shuffle=False, + num_workers=0, + ))) + _, peak = tracemalloc.get_traced_memory() + return peak + finally: + tracemalloc.stop() + + +def _repeat_permutations(size, count, seed): + values = [] + generator = np.random.RandomState(seed) + while len(values) < count: + values.extend(generator.permutation(size).tolist()) + return values[:count] + + +def _seed_everything(seed): + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + torch.use_deterministic_algorithms(True) + + +def _finite_float(value, name): + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError("%s must be scalar." % name) + value = float(value.detach()) + else: + value = float(value) + if not math.isfinite(value): + raise FloatingPointError("%s is NaN or Inf." % name) + return value + + +def _positive_int(value, name): + if isinstance(value, bool) or not isinstance(value, int) or value <= 0: + raise ValueError("%s must be a positive int." % name) + return value + + +class _SequenceDataset(Dataset): + def __init__(self, dataset, indices): + self._dataset = dataset + self._indices = indices + + def __len__(self): + return len(self._indices) + + def __getitem__(self, index): + return self._dataset[self._indices[index]] + + def __getitems__(self, indices): + source_indices = [self._indices[index] for index in indices] + getitems = getattr(self._dataset, "__getitems__", None) + if getitems is not None: + return getitems(source_indices) + return [self._dataset[index] for index in source_indices] diff --git a/paimon-python/pypaimon/benchmark/paired_act.py b/paimon-python/pypaimon/benchmark/paired_act.py new file mode 100644 index 000000000000..f36643aaa8e7 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/paired_act.py @@ -0,0 +1,835 @@ +# 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. + +"""Paired RoboMIND ACT benchmark over original HDF5 and Paimon. + +Both adapters consume one immutable :class:`BenchmarkConfig`, one train-only +normalization object, and one explicit window plan. The runner resets the same +seed before constructing the same LeRobot ACT policy and AdamW trainer for each +backend. It measures three alternating rounds without attempting OS cache +control and writes tensor and loss parity alongside timing and memory evidence. +Ingestion and canonical-action backfill are deliberately outside the benchmark. +""" + +import argparse +import gc +import hashlib +import json +import os +import platform +import subprocess +import time +import uuid +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + +import numpy as np +import torch +from torch.utils.data import Dataset + +import pypaimon.multimodal as pmm +from pypaimon.benchmark.act_harness import ( + BenchmarkConfig, + build_window_plan, + decode_rgb_image, + run_backend, +) +from pypaimon.sample import robomind_agilex as agilex + + +QPOS_COLUMNS = ( + "state_joint_position_left", + "state_joint_position_right", +) +ACTION_COLUMNS = ("action",) +IMAGE_COLUMNS = ( + "rgb_front", + "rgb_left_wrist", + "rgb_right_wrist", +) +HDF5_QPOS_FIELDS = ( + "puppet/joint_position_left", + "puppet/joint_position_right", +) +HDF5_ACTION_FIELDS = ( + "master/joint_position_left", + "master/joint_position_right", +) +HDF5_IMAGE_FIELDS = ( + "observations/rgb_images/camera_front", + "observations/rgb_images/camera_left_wrist", + "observations/rgb_images/camera_right_wrist", +) + + +@dataclass(frozen=True) +class _BenchmarkEpisode: + path: Path + source_key: str + episode_id: str + split: str + success: bool + frame_count: int + + +class Hdf5ACTWindowDataset(Dataset): + """Map-style complete ACT windows read on demand from one HDF5 episode.""" + + def __init__(self, episode, normalization, action_horizon): + self.episode = episode + self.normalization = normalization + self.action_horizon = action_horizon + self.window_count = episode.frame_count - action_horizon + 1 + if self.window_count <= 0: + raise ValueError( + "Episode %s is shorter than action horizon %d." + % (episode.episode_id, action_horizon)) + + def __len__(self): + return self.window_count + + def __getitem__(self, anchor): + if anchor < 0: + anchor += self.window_count + if anchor < 0 or anchor >= self.window_count: + raise IndexError(anchor) + import h5py + + with h5py.File(str(self.episode.path), "r") as h5: + qpos = _read_vectors(h5, HDF5_QPOS_FIELDS, anchor) + action = _read_vectors( + h5, + HDF5_ACTION_FIELDS, + slice(anchor, anchor + self.action_horizon), + ) + images = np.stack([ + _decode_hdf5_image(h5[field][anchor]) + for field in HDF5_IMAGE_FIELDS + ]) + qpos = ( + (qpos - self.normalization["qpos_mean"]) + / self.normalization["qpos_std"]) + action = ( + (action - self.normalization["action_mean"]) + / self.normalization["action_std"]) + return { + "sample_id": "%s#%d" % (self.episode.episode_id, anchor), + "episode_id": self.episode.episode_id, + "step_idx": anchor, + "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), + "action": torch.from_numpy(np.ascontiguousarray(action)), + "images": torch.from_numpy(np.ascontiguousarray(images)), + "is_pad": torch.zeros(self.action_horizon, dtype=torch.bool), + } + + +class _PaimonACTAdapter: + """Adapt one generic Paimon row window to the shared ACT contract.""" + + def __init__(self, normalization): + self.normalization = normalization + + def __call__(self, sample): + qpos = np.concatenate([ + np.asarray(sample[name][0], dtype=np.float32) + for name in QPOS_COLUMNS + ]) + action = np.concatenate([ + np.asarray(sample[name], dtype=np.float32) + for name in ACTION_COLUMNS + ], axis=-1) + images = np.stack([ + _decode_image(sample[name][0]) + for name in IMAGE_COLUMNS + ]) + qpos = ( + (qpos - self.normalization["qpos_mean"]) + / self.normalization["qpos_std"]) + action = ( + (action - self.normalization["action_mean"]) + / self.normalization["action_std"]) + episode_id = sample["episode_id"] + step_idx = sample["frame_index"] + return { + "sample_id": "%s#%d" % (episode_id, step_idx), + "episode_id": episode_id, + "step_idx": step_idx, + "qpos": torch.from_numpy(np.ascontiguousarray(qpos)), + "action": torch.from_numpy(np.ascontiguousarray(action)), + "images": torch.from_numpy(np.ascontiguousarray(images)), + "is_pad": sample["is_pad"], + } + + +def run( + input_root, + warehouse, + report_path, + *, + config=None, + database=agilex.DEFAULT_DATABASE, + statistics_version=agilex.DEFAULT_STATISTICS_VERSION, + train_episode_id=None, + validation_episode_id=None, + policy_factory=None): + """Run the paired benchmark without performing ingest or backfill.""" + config = config or BenchmarkConfig() + if not isinstance(config, BenchmarkConfig): + raise TypeError("config must be a BenchmarkConfig.") + started_at = _utc_now() + started = time.monotonic() + input_root = Path(input_root).expanduser().resolve() + warehouse = Path(warehouse).expanduser().resolve() + report_path = Path(report_path).expanduser().resolve() + + discovered_episodes = agilex.discover_episodes(input_root) + connection = pmm.connect( + database=database, options={"warehouse": str(warehouse)}) + episode_rows = _episode_rows(connection) + source_episodes, source_identity_sha256 = _validate_source_identity( + discovered_episodes, episode_rows) + source_by_id = {episode.episode_id: episode for episode in source_episodes} + frames = connection.get_table(agilex.FRAMES_TABLE) + frames_snapshot_id = _snapshot_id(frames) + + normalization, normalization_metadata = _shared_normalization( + source_episodes, + connection, + frames_snapshot_id, + statistics_version, + ) + train_episode = _select_episode( + source_by_id, + split="train", + requested=train_episode_id, + action_horizon=config.action_horizon, + ) + validation_episode = _select_episode( + source_by_id, + split="val", + requested=validation_episode_id, + action_horizon=config.action_horizon, + ) + plan = build_window_plan( + train_episode.frame_count - config.action_horizon + 1, + validation_episode.frame_count - config.action_horizon + 1, + config, + ) + sequence_sha256 = _sample_sequence_sha256( + train_episode.episode_id, + validation_episode.episode_id, + plan, + ) + + factories = { + "hdf5": lambda: _hdf5_datasets( + train_episode, validation_episode, normalization, config), + "paimon": lambda: _paimon_datasets( + frames, + frames_snapshot_id, + train_episode.episode_id, + validation_episode.episode_id, + normalization, + config, + ), + } + tensor_parity = _tensor_parity( + factories["hdf5"](), factories["paimon"](), plan) + del source_by_id + gc.collect() + + runs = [] + execution_order = [] + for round_index in range(config.rounds): + order = ( + ("hdf5", "paimon") + if round_index % 2 == 0 else ("paimon", "hdf5")) + for backend in order: + execution_order.append(backend) + runs.append(run_backend( + backend, + round_index + 1, + factories[backend], + plan, + config, + sequence_sha256, + policy_factory=policy_factory, + )) + gc.collect() + + loss_parity = _loss_parity(runs, config.rounds) + checks = { + "source_hdf5_matches_paimon": True, + "versioned_action_normalization_matches_hdf5": True, + "shared_normalization_object": True, + "shared_config": True, + "shared_seed": True, + "paimon_windows_snapshot_pinned": True, + "shared_window_sequence": len({ + item["sample_sequence_sha256"] for item in runs + }) == 1, + "tensor_parity": tensor_parity["passed"], + "train_and_validation_loss_parity": loss_parity["passed"], + "three_or_more_alternating_rounds": ( + config.rounds >= 3 + and execution_order == _expected_order(config.rounds)), + "all_losses_finite": all( + np.isfinite(value) + for item in runs + for value in item["train_loss"] + [item["validation_loss"]]), + } + status = "SUCCEEDED" if all(checks.values()) else "FAILED" + report = { + "schema_version": "robomind-paired-act-benchmark@1", + "benchmark_id": "M0-paired-ACT", + "run_id": "%s-%s" % ( + started_at.replace(":", "").replace("-", ""), + uuid.uuid4().hex[:8], + ), + "status": status, + "scope": "paired CPU ACT training path; ingest and backfill excluded", + "input": { + "dataset": "RoboMIND AgileX", + "input_manifest_sha256": source_identity_sha256, + "episode_count": len(source_episodes), + "warehouse": str(warehouse), + "database": database, + "frames_table": agilex.FRAMES_TABLE, + "frames_snapshot_id": frames_snapshot_id, + "paimon_window_dataset": ( + "pypaimon.multimodal.ContiguousWindowDataset"), + "paimon_window_snapshot_id": frames_snapshot_id, + "train_episode_id": train_episode.episode_id, + "validation_episode_id": validation_episode.episode_id, + }, + "parameters": { + "config": config.to_dict(), + "cache_control": "uncontrolled", + "device": "cpu", + "data_loader_workers": 0, + }, + "normalization": normalization_metadata, + "window_plan": { + **plan.to_dict(), + "sha256": plan.sha256, + "sample_sequence_sha256": sequence_sha256, + "train_episode_id": train_episode.episode_id, + "validation_episode_id": validation_episode.episode_id, + }, + "execution_order": execution_order, + "runs": runs, + "summary": { + backend: _summarize( + [item for item in runs if item["backend"] == backend]) + for backend in ("hdf5", "paimon") + }, + "correctness": { + "passed": all(checks.values()), + "checks": checks, + "tensor_parity": tensor_parity, + "loss_parity": loss_parity, + }, + "environment": { + "python": platform.python_version(), + "os": platform.platform(), + "machine": platform.machine(), + "torch": torch.__version__, + "source_commit": _git_head(Path(__file__).resolve().parents[3]), + }, + "command": _sanitized_command(), + "timing": {"wall_time_s": time.monotonic() - started}, + "unverified": [ + "OS page cache is uncontrolled; no cache dropping was attempted.", + "CPU fixed-step loss parity proves engineering equivalence, " + "not policy quality.", + "GPU, multi-worker DataLoader, distributed training, and " + "recovery are unverified.", + "Python tracemalloc does not include all native Arrow or " + "Torch allocations and is measured in a separate dataset-first-" + "batch replay.", + ], + "started_at": started_at, + "finished_at": _utc_now(), + } + if status != "SUCCEEDED": + raise AssertionError("Paired ACT correctness gate failed: %s" % checks) + report_path.parent.mkdir(parents=True, exist_ok=True) + report_path.write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return report + + +def _hdf5_datasets(train_episode, validation_episode, normalization, config): + return ( + Hdf5ACTWindowDataset( + train_episode, normalization, config.action_horizon), + Hdf5ACTWindowDataset( + validation_episode, normalization, config.action_horizon), + ) + + +def _paimon_datasets( + frames, + frames_snapshot_id, + train_episode_id, + validation_episode_id, + normalization, + config): + datasets = tuple( + frames.scan(snapshot_id=frames_snapshot_id).where( + "episode_id = '%s'" % episode_id.replace("'", "''") + ).to_contiguous_window_dataset( + window_size=config.action_horizon, + columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, + group_key="episode_id", + order_key="frame_index", + stride=1, + tail="drop", + adapter=_PaimonACTAdapter(normalization), + ) + for episode_id in (train_episode_id, validation_episode_id) + ) + actual_snapshot_ids = {dataset.snapshot_id for dataset in datasets} + if actual_snapshot_ids != {frames_snapshot_id}: + raise RuntimeError( + "Paimon ACT windows must remain pinned to frames snapshot %s; " + "got %s." + % (frames_snapshot_id, sorted(actual_snapshot_ids))) + return datasets + + +def _shared_normalization( + episodes, + connection, + frames_snapshot_id, + statistics_version): + train = [ + episode for episode in episodes + if episode.split == "train" and episode.success + ] + if not train: + raise ValueError("No successful train episodes are available.") + qpos = _Moments(14) + action = _Moments(14) + import h5py + + for episode in sorted(train, key=lambda item: item.episode_id): + with h5py.File(str(episode.path), "r") as h5: + qpos.update(_read_vectors( + h5, HDF5_QPOS_FIELDS, slice(None), dtype=np.float64)) + action.update(_read_vectors( + h5, HDF5_ACTION_FIELDS, slice(None), dtype=np.float64)) + qpos_mean, qpos_std = qpos.finish() + action_mean, action_std = action.finish() + row = _statistics_row(connection, statistics_version) + if row["source_snapshot_id"] != frames_snapshot_id: + raise ValueError( + "Normalization source snapshot %s differs from frames " + "snapshot %s." + % (row["source_snapshot_id"], frames_snapshot_id)) + if row["source_split"] != "train" or row["frame_count"] != action.count: + raise ValueError( + "Versioned action normalization has the wrong train scope.") + if row["feature_name"] != "action": + raise ValueError("Versioned normalization feature must be action.") + if row["standard_deviation_floor"] != 1e-2: + raise ValueError( + "Versioned normalization must use the 1e-2 std floor.") + stored_mean = np.asarray(row["action_mean"], dtype=np.float64) + stored_std = np.asarray(row["action_std"], dtype=np.float64) + if not ( + np.allclose(stored_mean, action_mean, rtol=1e-10, atol=1e-10) + and np.allclose(stored_std, action_std, rtol=1e-10, atol=1e-10)): + raise ValueError( + "Versioned Paimon action normalization differs from HDF5 source.") + normalization = { + "qpos_mean": qpos_mean.astype(np.float32), + "qpos_std": qpos_std.astype(np.float32), + "action_mean": stored_mean.astype(np.float32), + "action_std": stored_std.astype(np.float32), + } + serializable = { + name: value.tolist() for name, value in normalization.items() + } + digest = hashlib.sha256(json.dumps( + serializable, sort_keys=True, separators=(",", ":") + ).encode("utf-8")).hexdigest() + return normalization, { + "statistics_version": statistics_version, + "source_split": "train", + "frame_count": action.count, + "standard_deviation_floor": 1e-2, + "values": serializable, + "sha256": digest, + } + + +def _statistics_row(connection, statistics_version): + escaped = statistics_version.replace("'", "''") + rows = (connection.get_table(agilex.FEATURE_STATS_TABLE).scan() + .where("statistics_version = '%s'" % escaped).to_list()) + if len(rows) != 1: + raise ValueError( + "Expected one normalization row for %r, got %d." + % (statistics_version, len(rows))) + return rows[0] + + +def _episode_rows(connection): + return connection.get_table(agilex.EPISODES_TABLE).scan().select([ + "episode_id", + "source_key", + "split", + "success", + "frame_count", + ]).to_list() + + +def _validate_source_identity(episodes, rows): + expected = { + item.episode_id: { + "episode_id": item.episode_id, + "source_key": item.source_key, + "split": item.split, + "success": item.success, + } + for item in episodes + } + actual = { + item["episode_id"]: { + "episode_id": item["episode_id"], + "source_key": item["source_key"], + "split": item["split"], + "success": item["success"], + } + for item in rows + } + if actual != expected or len(actual) != len(rows): + raise ValueError( + "HDF5 and Paimon source identity differ; rebuild or select " + "matching inputs.") + rows_by_id = {item["episode_id"]: item for item in rows} + enriched = [ + _BenchmarkEpisode( + path=item.path, + source_key=item.source_key, + episode_id=item.episode_id, + split=item.split, + success=item.success, + frame_count=rows_by_id[item.episode_id]["frame_count"], + ) + for item in episodes + ] + manifest = sorted([ + { + "episode_id": item.episode_id, + "source_key": item.source_key, + "split": item.split, + "success": item.success, + "frame_count": rows_by_id[item.episode_id]["frame_count"], + } + for item in episodes + ], key=lambda item: item["episode_id"]) + payload = json.dumps(manifest, sort_keys=True, separators=(",", ":")) + return enriched, hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def _select_episode(source_by_id, split, requested, action_horizon): + eligible = { + episode_id: episode + for episode_id, episode in source_by_id.items() + if episode.split == split + and episode.success + and episode.frame_count >= action_horizon + } + if not eligible: + raise ValueError( + "No successful %s episode is long enough for horizon %d." + % (split, action_horizon)) + selected = requested or min(eligible) + if selected not in eligible: + raise ValueError( + "Requested %s episode is missing, unsuccessful, or too short: %s." + % (split, selected)) + return eligible[selected] + + +def _tensor_parity(hdf5_datasets, paimon_datasets, plan): + comparisons = ( + ("train", hdf5_datasets[0], paimon_datasets[0], + sorted(set(plan.loader_indices + plan.train_indices))), + ("validation", hdf5_datasets[1], paimon_datasets[1], + sorted(set(plan.validation_indices))), + ) + checked = 0 + max_absolute_difference = { + "qpos": 0.0, + "action": 0.0, + "images": 0.0, + } + for split, hdf5_dataset, paimon_dataset, indices in comparisons: + if len(hdf5_dataset) != len(paimon_dataset): + raise AssertionError( + "%s window counts differ: HDF5=%d Paimon=%d." + % (split, len(hdf5_dataset), len(paimon_dataset))) + for index in indices: + hdf5_sample = hdf5_dataset[index] + paimon_sample = paimon_dataset[index] + for name in ("sample_id", "episode_id", "step_idx"): + if hdf5_sample[name] != paimon_sample[name]: + raise AssertionError( + "%s %s differs at window %d." % (split, name, index)) + for name in ("qpos", "action", "images", "is_pad"): + if not torch.equal(hdf5_sample[name], paimon_sample[name]): + raise AssertionError( + "%s %s tensor differs at %s." + % (split, name, hdf5_sample["sample_id"])) + if name in max_absolute_difference: + difference = torch.max(torch.abs( + hdf5_sample[name] - paimon_sample[name])).item() + max_absolute_difference[name] = max( + max_absolute_difference[name], difference) + checked += 1 + return { + "passed": True, + "checked_window_count": checked, + "comparison": "torch.equal", + "max_absolute_difference": max_absolute_difference, + } + + +def _loss_parity(runs, round_count): + comparisons = [] + passed = True + for round_number in range(1, round_count + 1): + by_backend = { + item["backend"]: item + for item in runs if item["round"] == round_number + } + hdf5_train = np.asarray(by_backend["hdf5"]["train_loss"]) + paimon_train = np.asarray(by_backend["paimon"]["train_loss"]) + train_equal = np.array_equal(hdf5_train, paimon_train) + validation_equal = ( + by_backend["hdf5"]["validation_loss"] + == by_backend["paimon"]["validation_loss"]) + passed = passed and train_equal and validation_equal + comparisons.append({ + "round": round_number, + "train_loss_exact": bool(train_equal), + "validation_loss_exact": bool(validation_equal), + "train_max_absolute_difference": float(np.max(np.abs( + hdf5_train - paimon_train))), + "validation_absolute_difference": abs( + by_backend["hdf5"]["validation_loss"] + - by_backend["paimon"]["validation_loss"]), + }) + return { + "passed": bool(passed), + "comparison": "exact CPU deterministic equality", + "rounds": comparisons, + } + + +def _summarize(runs): + metrics = ( + "dataset_build_s", + "first_batch_s", + "dataloader_samples_per_s", + "fixed_steps_s", + "validation_loss", + "python_peak_allocated_bytes", + "wall_time_s", + ) + result = {"round_count": len(runs)} + for name in metrics: + values = [item[name] for item in runs] + result[name] = { + "median": float(np.median(values)), + "min": float(np.min(values)), + "max": float(np.max(values)), + } + return result + + +def _sample_sequence_sha256(train_episode_id, validation_episode_id, plan): + value = { + "loader": [ + "%s#%d" % (train_episode_id, index) + for index in plan.loader_indices + ], + "train": [ + "%s#%d" % (train_episode_id, index) + for index in plan.train_indices + ], + "validation": [ + "%s#%d" % (validation_episode_id, index) + for index in plan.validation_indices + ], + } + return hashlib.sha256(json.dumps( + value, sort_keys=True, separators=(",", ":") + ).encode("utf-8")).hexdigest() + + +def _read_vectors(h5, fields, selection, dtype=np.float32): + value = np.concatenate([ + np.asarray(h5[field][selection], dtype=dtype) + for field in fields + ], axis=-1) + if not np.isfinite(value).all(): + raise ValueError("ACT vector contains NaN or Inf.") + return value + + +def _decode_hdf5_image(value): + return _decode_image(value) + + +def _decode_image(value): + payload = ( + bytes(value) + if isinstance(value, (bytes, bytearray, memoryview)) + else np.asarray(value, dtype=np.uint8).tobytes() + ) + image = decode_rgb_image(payload) + return np.transpose(image, (2, 0, 1)).astype(np.float32) / 255.0 + + +class _Moments(object): + def __init__(self, width): + self.count = 0 + self.total = np.zeros(width, dtype=np.float64) + self.total_square = np.zeros(width, dtype=np.float64) + + def update(self, value): + value = np.asarray(value, dtype=np.float64) + if value.ndim != 2 or value.shape[1] != len(self.total): + raise ValueError( + "Unexpected normalization shape %s." % (value.shape,)) + if not np.isfinite(value).all(): + raise ValueError("Normalization input contains NaN or Inf.") + self.count += value.shape[0] + self.total += value.sum(axis=0) + self.total_square += np.square(value).sum(axis=0) + + def finish(self): + if self.count == 0: + raise ValueError("Cannot compute normalization from no frames.") + mean = self.total / self.count + variance = np.maximum( + self.total_square / self.count - np.square(mean), 0.0) + return mean, np.maximum(np.sqrt(variance), 1e-2) + + +def _snapshot_id(table): + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() + if snapshot is None: + raise ValueError("Paimon frames table has no snapshot.") + return snapshot.id + + +def _expected_order(rounds): + result = [] + for index in range(rounds): + result.extend( + ("hdf5", "paimon") if index % 2 == 0 else ("paimon", "hdf5")) + return result + + +def _git_head(repository): + return subprocess.check_output( + ["git", "-C", str(repository), "rev-parse", "HEAD"], + universal_newlines=True, + ).strip() + + +def _sanitized_command(): + import sys + return [os.path.basename(sys.executable)] + list(sys.argv) + + +def _utc_now(): + return datetime.now(timezone.utc).isoformat( + timespec="seconds").replace("+00:00", "Z") + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", required=True) + parser.add_argument("--warehouse", required=True) + parser.add_argument("--report", required=True) + parser.add_argument("--database", default=agilex.DEFAULT_DATABASE) + parser.add_argument( + "--statistics-version", default=agilex.DEFAULT_STATISTICS_VERSION) + parser.add_argument("--train-episode-id") + parser.add_argument("--validation-episode-id") + parser.add_argument("--seed", type=int, default=BenchmarkConfig.seed) + parser.add_argument( + "--action-horizon", type=int, default=BenchmarkConfig.action_horizon) + parser.add_argument( + "--batch-size", type=int, default=BenchmarkConfig.batch_size) + parser.add_argument( + "--optimizer-steps", type=int, default=BenchmarkConfig.optimizer_steps) + parser.add_argument( + "--image-height", type=int, default=BenchmarkConfig.image_height) + parser.add_argument( + "--image-width", type=int, default=BenchmarkConfig.image_width) + parser.add_argument( + "--learning-rate", type=float, default=BenchmarkConfig.learning_rate) + parser.add_argument( + "--weight-decay", type=float, default=BenchmarkConfig.weight_decay) + parser.add_argument( + "--warmup-batches", type=int, default=BenchmarkConfig.warmup_batches) + parser.add_argument( + "--loader-batches", type=int, default=BenchmarkConfig.loader_batches) + parser.add_argument("--rounds", type=int, default=BenchmarkConfig.rounds) + args = parser.parse_args(argv) + config = BenchmarkConfig( + seed=args.seed, + action_horizon=args.action_horizon, + batch_size=args.batch_size, + optimizer_steps=args.optimizer_steps, + image_height=args.image_height, + image_width=args.image_width, + learning_rate=args.learning_rate, + weight_decay=args.weight_decay, + warmup_batches=args.warmup_batches, + loader_batches=args.loader_batches, + rounds=args.rounds, + ) + report = run( + args.input, + args.warehouse, + args.report, + config=config, + database=args.database, + statistics_version=args.statistics_version, + train_episode_id=args.train_episode_id, + validation_episode_id=args.validation_episode_id, + ) + print(json.dumps({ + "status": report["status"], + "report": str(Path(args.report).expanduser().resolve()), + "input_manifest_sha256": report["input"]["input_manifest_sha256"], + }, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 6517d9e1d1e4..d47f2c79c895 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -115,16 +115,41 @@ def __len__(self): return len(self._anchors) def __getitem__(self, index): + anchor, row_ids = self._resolve_window(index) + return self._sample(anchor, self._read_rows(row_ids)) + + def __getitems__(self, indices): + windows = [self._resolve_window(index) for index in indices] + if not windows: + return [] + row_ids = list(dict.fromkeys( + row_id for _, window_row_ids in windows + for row_id in window_row_ids + )) + rows_by_id = dict(zip(row_ids, self._read_rows(row_ids))) + return [ + self._sample( + anchor, + [rows_by_id[row_id] for row_id in window_row_ids], + ) + for anchor, window_row_ids in windows + ] + + def _resolve_window(self, index): index = operator.index(index) if index < 0: index += len(self._anchors) if index < 0 or index >= len(self._anchors): raise IndexError("window index out of range") - group_index, start, valid_count = self._anchors[index] - group_key, order_values, row_ids = self._groups[group_index] - selected_row_ids = row_ids[start:start + valid_count] - rows = self._read_rows(selected_row_ids) + anchor = self._anchors[index] + group_index, start, valid_count = anchor + row_ids = self._groups[group_index][2] + return anchor, row_ids[start:start + valid_count] + + def _sample(self, anchor, rows): + group_index, start, valid_count = anchor + group_key, order_values, _ = self._groups[group_index] padding_count = self.window_size - valid_count padding_mask = torch.zeros(self.window_size, dtype=torch.bool) if padding_count: diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index 34ce7e647d1a..dc7112a6800b 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -145,6 +145,26 @@ def test_reads_blob_payloads_only_when_a_window_is_requested(self): sample["payload"], ) + def test_plural_access_coalesces_overlapping_window_reads(self): + dataset = self._dataset(self._table()) + expected = [dataset[0], dataset[1]] + + with patch.object( + dataset, "_read_rows", wraps=dataset._read_rows) as read: + actual = dataset.__getitems__([0, 1]) + + self.assertEqual(1, read.call_count) + self.assertEqual(4, len(read.call_args.args[0])) + for expected_sample, actual_sample in zip(expected, actual): + self.assertEqual( + expected_sample["episode"], actual_sample["episode"]) + self.assertEqual(expected_sample["step"], actual_sample["step"]) + self.assertEqual(expected_sample["value"], actual_sample["value"]) + self.assertEqual( + expected_sample["payload"], actual_sample["payload"]) + self.assertTrue(torch.equal( + expected_sample["is_pad"], actual_sample["is_pad"])) + def test_pad_tail_repeats_last_row_and_marks_real_padding(self): dataset = self._dataset( self._table(), tail="pad", pad_values={"value": -1}) diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py new file mode 100644 index 000000000000..40e302f9765a --- /dev/null +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -0,0 +1,343 @@ +# 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 json +import tracemalloc +from io import BytesIO +from unittest.mock import patch + +import numpy as np +import pytest +import torch +from PIL import Image + +import pypaimon.multimodal as pmm +from pypaimon.benchmark.paired_act import ( + IMAGE_COLUMNS, + BenchmarkConfig, + _paimon_datasets, + _shared_normalization, + _snapshot_id, + run, +) +from pypaimon.benchmark.act_harness import ( + _SequenceDataset, + build_window_plan, + run_backend, +) +from pypaimon.multimodal.query import ScanQuery +from pypaimon.multimodal.window_dataset import ContiguousWindowDataset +from pypaimon.sample import robomind_agilex as agilex + + +h5py = pytest.importorskip("h5py") + + +def test_sequence_dataset_forwards_plural_access(): + class BatchDataset: + def __getitems__(self, indices): + return ["sample-%d" % index for index in indices] + + dataset = _SequenceDataset(BatchDataset(), (7, 3, 5)) + + assert dataset.__getitems__([0, 2]) == ["sample-7", "sample-5"] + + +def test_backend_times_without_tracemalloc_and_measures_memory_separately(): + states = [] + config = BenchmarkConfig( + seed=11, + action_horizon=1, + batch_size=1, + optimizer_steps=1, + image_height=2, + image_width=2, + warmup_batches=1, + loader_batches=1, + rounds=3, + ) + + class TracingDataset(torch.utils.data.Dataset): + def __len__(self): + return 2 + + def __getitem__(self, index): + states.append(tracemalloc.is_tracing()) + return { + "sample_id": "episode-a#%d" % index, + "episode_id": "episode-a", + "step_idx": index, + "qpos": torch.zeros(14), + "action": torch.zeros((1, 14)), + "images": torch.zeros((3, 3, 2, 2)), + "is_pad": torch.zeros(1, dtype=torch.bool), + } + + dataset = TracingDataset() + plan = build_window_plan(len(dataset), len(dataset), config) + result = run_backend( + "test", + 1, + lambda: (dataset, dataset), + plan, + config, + "sequence-sha256", + policy_factory=_policy_factory, + ) + + assert states[0] is False + assert states[-1] is True + assert result["peak_memory_measurement"] == ( + "python-tracemalloc-separate-dataset-first-batch") + + +def _jpeg(value): + buffer = BytesIO() + Image.fromarray(np.full((8, 10, 3), value, dtype=np.uint8)).save( + buffer, format="JPEG") + return np.frombuffer(buffer.getvalue(), dtype=np.uint8) + + +def _write_episode(root, split, name, offset, frames=6): + path = (root / "13_packbowl" / "success_episodes" / split / name + / "data" / "trajectory.hdf5") + path.parent.mkdir(parents=True) + with h5py.File(path, "w") as h5: + h5.create_dataset("language_raw", data=[b"pack the bowl"]) + h5.create_dataset( + "language_distilbert", + data=np.zeros((1, 1, 768), dtype=np.float16), + ) + for index, (_, hdf5_path) in enumerate(agilex.NUMERIC_FIELDS): + values = np.arange(frames * 7, dtype=np.float64).reshape(frames, 7) + h5.create_dataset(hdf5_path, data=values + offset + index * 100) + variable = h5py.vlen_dtype(np.dtype("uint8")) + for image_index, (_, hdf5_path) in enumerate(agilex.IMAGE_FIELDS): + dataset = h5.create_dataset(hdf5_path, (frames,), dtype=variable) + for frame_index in range(frames): + dataset[frame_index] = _jpeg( + offset + image_index + frame_index) + return path + + +@pytest.fixture +def paired_input(tmp_path): + root = tmp_path / "input" + _write_episode(root, "train", "train-a", 1) + _write_episode(root, "train", "train-b", 11) + _write_episode(root, "val", "val-a", 21) + warehouse = tmp_path / "warehouse" + agilex.ingest_local(root, warehouse, batch_size=2) + agilex.backfill_canonical_action( + warehouse, statistics_version="paired-test@1") + return root, warehouse + + +class _Policy(torch.nn.Module): + + def __init__(self): + super().__init__() + self.scale = torch.nn.Parameter(torch.tensor(0.0)) + + def forward(self, batch): + assert self.training + target = batch["action"].mean() + batch["observation.state"].mean() + loss = (self.scale - target).square() + return loss, { + "l1_loss": loss.detach(), + "kld_loss": torch.tensor(0.0), + } + + +def _policy_factory(config): + return _Policy(), { + "implementation": "test-policy", + "chunk_size": config.action_horizon, + "parameter_count": 1, + } + + +def test_runs_three_alternating_rounds_with_one_shared_contract( + paired_input, tmp_path): + input_root, warehouse = paired_input + report_path = tmp_path / "paired-report.json" + config = BenchmarkConfig( + seed=17, + action_horizon=3, + batch_size=2, + optimizer_steps=2, + image_height=8, + image_width=10, + warmup_batches=1, + loader_batches=2, + rounds=3, + ) + + report = run( + input_root, + warehouse, + report_path, + config=config, + statistics_version="paired-test@1", + policy_factory=_policy_factory, + ) + + assert report_path.exists() + assert json.loads(report_path.read_text()) == report + assert report["schema_version"] == "robomind-paired-act-benchmark@1" + assert report["status"] == "SUCCEEDED" + assert report["parameters"]["config"] == config.to_dict() + assert report["parameters"]["cache_control"] == "uncontrolled" + assert report["input"]["paimon_window_dataset"] == ( + "pypaimon.multimodal.ContiguousWindowDataset") + assert report["input"]["paimon_window_snapshot_id"] == ( + report["input"]["frames_snapshot_id"]) + assert report["execution_order"] == [ + "hdf5", "paimon", "paimon", "hdf5", "hdf5", "paimon", + ] + assert len(report["runs"]) == 6 + assert all(report["correctness"]["checks"].values()) + assert report["correctness"]["tensor_parity"]["passed"] + assert report["correctness"]["tensor_parity"][ + "checked_window_count"] > 0 + assert ( + report["correctness"]["tensor_parity"]["max_absolute_difference"] + == { + "qpos": 0.0, + "action": 0.0, + "images": 0.0, + } + ) + assert report["correctness"]["loss_parity"]["passed"] + assert all( + comparison["train_loss_exact"] + and comparison["validation_loss_exact"] + for comparison in report["correctness"]["loss_parity"]["rounds"] + ) + assert report["window_plan"]["seed"] == 17 + assert len(report["window_plan"]["sha256"]) == 64 + assert report["normalization"]["statistics_version"] == "paired-test@1" + assert len(report["normalization"]["sha256"]) == 64 + assert set(report["summary"]) == {"hdf5", "paimon"} + for backend in ("hdf5", "paimon"): + assert report["summary"][backend]["round_count"] == 3 + for metric in ( + "first_batch_s", + "dataloader_samples_per_s", + "fixed_steps_s", + "python_peak_allocated_bytes"): + assert set(report["summary"][backend][metric]) == { + "median", "min", "max", + } + for round_index in range(3): + paired = [item for item in report["runs"] + if item["round"] == round_index + 1] + by_backend = {item["backend"]: item for item in paired} + assert by_backend["hdf5"]["sample_sequence_sha256"] == ( + by_backend["paimon"]["sample_sequence_sha256"]) + assert by_backend["hdf5"]["train_loss"] == ( + by_backend["paimon"]["train_loss"]) + assert by_backend["hdf5"]["validation_loss"] == ( + by_backend["paimon"]["validation_loss"]) + + +def test_paimon_windows_are_lazy_and_snapshot_pinned(paired_input): + input_root, warehouse = paired_input + connection = pmm.connect( + database=agilex.DEFAULT_DATABASE, + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(agilex.FRAMES_TABLE) + snapshot_id = _snapshot_id(frames) + normalization, _ = _shared_normalization( + agilex.discover_episodes(input_root), + connection, + snapshot_id, + "paired-test@1", + ) + + original = ScanQuery._fetch_bodies + with patch.object( + ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + train, validation = _paimon_datasets( + frames, + snapshot_id, + "train-a", + "val-a", + normalization, + BenchmarkConfig( + action_horizon=3, + batch_size=1, + optimizer_steps=1, + image_height=8, + image_width=10, + rounds=3, + ), + ) + assert fetch.call_count == 0 + assert isinstance(train, ContiguousWindowDataset) + assert isinstance(validation, ContiguousWindowDataset) + sample_before_append = train[0] + assert fetch.call_count == 1 + + scalar, blobs = frames.scan().where( + "episode_id = 'train-a' AND frame_index = 5" + ).read_blobs(IMAGE_COLUMNS) + appended = scalar.to_pylist()[0] + appended["frame_index"] = 6 + for name in IMAGE_COLUMNS: + appended[name] = blobs[name][0] + frames.add([appended]) + + assert train.snapshot_id == snapshot_id + assert validation.snapshot_id == snapshot_id + assert _snapshot_id(frames) != snapshot_id + assert len(train) == 4 + sample_after_append = train[0] + for name in ("qpos", "action", "images", "is_pad"): + assert torch.equal( + sample_before_append[name], sample_after_append[name]) + + +def test_tensor_parity_rejects_different_hdf5_bytes( + paired_input, tmp_path): + input_root, warehouse = paired_input + changed = (input_root / "13_packbowl" / "success_episodes" / "train" + / "train-a" / "data" / "trajectory.hdf5") + with h5py.File(changed, "r+") as h5: + h5["puppet/joint_position_left"][0, 0] += 1 + + with pytest.raises(AssertionError, match="tensor differs"): + run( + input_root, + warehouse, + tmp_path / "must-not-exist.json", + config=BenchmarkConfig( + action_horizon=3, + batch_size=1, + optimizer_steps=1, + image_height=8, + image_width=10, + rounds=3, + ), + statistics_version="paired-test@1", + policy_factory=_policy_factory, + ) + + +def test_requires_at_least_three_alternating_rounds(): + with pytest.raises(ValueError, match="rounds must be at least 3"): + BenchmarkConfig(rounds=2) diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 9111730fb7ca..cccd2838e24b 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -243,6 +243,10 @@ def read_requirements(): 'torch': [ 'torch', ], + 'act': [ + 'lerobot==0.4.4', + 'Pillow', + ], 'daft': [ 'daft>=0.7.6; python_version>="3.10"', ], From 5d8890ee890f2d1d66c5d77316a0ec791db0a24f Mon Sep 17 00:00:00 2001 From: Yann Date: Sat, 29 Aug 2026 22:43:07 +0800 Subject: [PATCH 3/3] fix(python): harden paired ACT benchmark Read observation images only at each window anchor and tolerate installed packages without a Git checkout. Improve public option documentation and CLI guidance. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 200/200 AI-Contributed/UT: 25/25 --- docs/docs/pypaimon/pytorch.md | 5 +- docs/docs/pypaimon/robomind-act-benchmark.md | 15 ++- .../pypaimon/benchmark/act_harness.py | 4 +- .../pypaimon/benchmark/paired_act.py | 76 ++++++++++----- paimon-python/pypaimon/multimodal/query.py | 7 +- .../pypaimon/multimodal/window_dataset.py | 93 +++++++++++++++---- .../tests/contiguous_window_dataset_test.py | 13 +++ .../tests/paired_act_benchmark_test.py | 12 +++ 8 files changed, 175 insertions(+), 50 deletions(-) diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index 2e6e3cfc9994..26020d72fe09 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -122,6 +122,7 @@ dataset = ( .to_contiguous_window_dataset( window_size=16, columns=["state", "image"], + anchor_columns=["image"], group_key="episode_id", order_key="step_idx", tail="pad", @@ -134,7 +135,9 @@ loader = DataLoader(dataset, batch_size=32, num_workers=4, shuffle=True) Each item contains the group and order keys, one list for each requested column, and a boolean `is_pad` tensor where `True` marks padding. Padding repeats the final real value by default; `pad_values` can override individual -columns. Use `column_transforms` to convert column lists to tensors and +columns. Columns named in `anchor_columns` contain only the first row's value, +which is useful when an observation applies to a full action window. Use +`column_transforms` to convert column lists to tensors and `adapter` to produce a model-specific sample mapping. Keep these callbacks picklable when using multiple DataLoader workers. diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index 0884553b9307..664465fc111e 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -37,20 +37,27 @@ python -m pypaimon.benchmark.paired_act \ --report /data/results/paired-act.json ``` +A successful run prints a compact `SUCCEEDED` result and writes the full JSON +report. A source, parity, or configuration mismatch raises an error and does +not write a successful report. + One immutable configuration controls both paths. The runner computes train-only normalization once, verifies its canonical action values against the requested version in `feature_stats_agilex`, and passes the same object to both adapters. A seeded window plan fixes every warmup, loader, training, and validation -anchor. Before training, the runner requires exact `torch.equal` parity for -sample identity, state, action, image, and padding tensors. +anchor. Before training, the runner compares `sample_id`, `episode_id`, and +`step_idx` by value and requires exact `torch.equal` parity for state, action, +image, and padding tensors. The Paimon adapter uses `ContiguousWindowDataset`, not an ACT-specific table reader. Dataset construction indexes only episode, frame, and row IDs. Window payloads remain lazy until `__getitem__`, and all train and validation reads are pinned to the exact frames snapshot recorded by the normalization statistics. PyTorch batch access coalesces overlapping row IDs into one payload read. The -adapter maps each generic window to the same tensor contract as the HDF5 adapter -without materializing episodes in memory. +image columns are marked as anchor-only, so each sample loads the observation +images once rather than once per action-horizon row. The adapter maps each +generic window to the same tensor contract as the HDF5 adapter without +materializing episodes in memory. Each backend then uses the same CPU LeRobot ACT policy, initial seed, AdamW optimizer, batch size, window sequence, and optimizer step count. At least diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act_harness.py index 15e5a7d3364c..dd874406579c 100644 --- a/paimon-python/pypaimon/benchmark/act_harness.py +++ b/paimon-python/pypaimon/benchmark/act_harness.py @@ -179,7 +179,7 @@ def validate_act_batch(batch, config): if torch.any(batch["images"] < 0) or torch.any(batch["images"] > 1): raise ValueError("images must be normalized to [0, 1].") if batch["is_pad"].any(): - raise ValueError("M0 ACT windows must be complete and unpadded.") + raise ValueError("Paired ACT benchmark windows must be complete and unpadded.") for sample_id, episode_id, step_idx in zip( batch["sample_id"], batch["episode_id"], batch["step_idx"].tolist()): @@ -209,7 +209,7 @@ def build_lerobot_batch(batch, config): def build_act_policy(config): - """Build the one reduced CPU LeRobot ACT configuration used by M0.""" + """Build the reduced CPU LeRobot ACT configuration used by the benchmark.""" try: import importlib.metadata from lerobot.configs.types import FeatureType, PolicyFeature diff --git a/paimon-python/pypaimon/benchmark/paired_act.py b/paimon-python/pypaimon/benchmark/paired_act.py index f36643aaa8e7..93cfa512d05b 100644 --- a/paimon-python/pypaimon/benchmark/paired_act.py +++ b/paimon-python/pypaimon/benchmark/paired_act.py @@ -397,6 +397,7 @@ def _paimon_datasets( ).to_contiguous_window_dataset( window_size=config.action_horizon, columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, + anchor_columns=IMAGE_COLUMNS, group_key="episode_id", order_key="frame_index", stride=1, @@ -753,10 +754,14 @@ def _expected_order(rounds): def _git_head(repository): - return subprocess.check_output( - ["git", "-C", str(repository), "rev-parse", "HEAD"], - universal_newlines=True, - ).strip() + try: + return subprocess.check_output( + ["git", "-C", str(repository), "rev-parse", "HEAD"], + stderr=subprocess.DEVNULL, + universal_newlines=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + return "UNKNOWN" def _sanitized_command(): @@ -770,35 +775,60 @@ def _utc_now(): def main(argv=None): - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--input", required=True) - parser.add_argument("--warehouse", required=True) - parser.add_argument("--report", required=True) - parser.add_argument("--database", default=agilex.DEFAULT_DATABASE) + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + parser.add_argument( + "--input", required=True, help="RoboMIND AgileX HDF5 root directory.") + parser.add_argument( + "--warehouse", required=True, help="Existing Paimon warehouse path.") + parser.add_argument( + "--report", required=True, help="Destination JSON report path.") + parser.add_argument( + "--database", default=agilex.DEFAULT_DATABASE, + help="Paimon database containing the ingested dataset.") + parser.add_argument( + "--statistics-version", default=agilex.DEFAULT_STATISTICS_VERSION, + help="Canonical action statistics version to verify and use.") + parser.add_argument( + "--train-episode-id", help="Train episode; defaults to the first eligible episode.") + parser.add_argument( + "--validation-episode-id", + help="Validation episode; defaults to the first eligible episode.") + parser.add_argument( + "--seed", type=int, default=BenchmarkConfig.seed, + help="Shared random seed and window-plan seed.") parser.add_argument( - "--statistics-version", default=agilex.DEFAULT_STATISTICS_VERSION) - parser.add_argument("--train-episode-id") - parser.add_argument("--validation-episode-id") - parser.add_argument("--seed", type=int, default=BenchmarkConfig.seed) + "--action-horizon", type=int, default=BenchmarkConfig.action_horizon, + help="Number of contiguous action rows in each sample.") parser.add_argument( - "--action-horizon", type=int, default=BenchmarkConfig.action_horizon) + "--batch-size", type=int, default=BenchmarkConfig.batch_size, + help="Shared DataLoader batch size.") parser.add_argument( - "--batch-size", type=int, default=BenchmarkConfig.batch_size) + "--optimizer-steps", type=int, default=BenchmarkConfig.optimizer_steps, + help="Fixed optimizer steps per backend run.") parser.add_argument( - "--optimizer-steps", type=int, default=BenchmarkConfig.optimizer_steps) + "--image-height", type=int, default=BenchmarkConfig.image_height, + help="ACT input image height after resizing.") parser.add_argument( - "--image-height", type=int, default=BenchmarkConfig.image_height) + "--image-width", type=int, default=BenchmarkConfig.image_width, + help="ACT input image width after resizing.") parser.add_argument( - "--image-width", type=int, default=BenchmarkConfig.image_width) + "--learning-rate", type=float, default=BenchmarkConfig.learning_rate, + help="Shared AdamW learning rate.") parser.add_argument( - "--learning-rate", type=float, default=BenchmarkConfig.learning_rate) + "--weight-decay", type=float, default=BenchmarkConfig.weight_decay, + help="Shared AdamW weight decay.") parser.add_argument( - "--weight-decay", type=float, default=BenchmarkConfig.weight_decay) + "--warmup-batches", type=int, default=BenchmarkConfig.warmup_batches, + help="DataLoader batches consumed before timing.") parser.add_argument( - "--warmup-batches", type=int, default=BenchmarkConfig.warmup_batches) + "--loader-batches", type=int, default=BenchmarkConfig.loader_batches, + help="Batches used for DataLoader throughput measurement.") parser.add_argument( - "--loader-batches", type=int, default=BenchmarkConfig.loader_batches) - parser.add_argument("--rounds", type=int, default=BenchmarkConfig.rounds) + "--rounds", type=int, default=BenchmarkConfig.rounds, + help="Alternating backend rounds; must be at least three.") args = parser.parse_args(argv) config = BenchmarkConfig( seed=args.seed, diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 4b614de11c49..310c0c1e11f0 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -156,6 +156,7 @@ def to_contiguous_window_dataset( *, window_size, columns=None, + anchor_columns=None, group_key="episode_id", order_key="step_idx", stride=1, @@ -167,8 +168,9 @@ def to_contiguous_window_dataset( """Build a snapshot-pinned, map-style Dataset of contiguous rows. The Dataset indexes only ``group_key``, ``order_key``, and Paimon row - IDs, then reads projected values on demand. It sorts rows within each - group and never creates a window across groups. See + IDs, then reads projected values on demand. Columns listed in + ``anchor_columns`` are read only for the first row of each window. It + sorts rows within each group and never creates a window across groups. See :class:`pypaimon.multimodal.window_dataset.ContiguousWindowDataset` for tail, padding, mask, transform, and adapter semantics. """ @@ -181,6 +183,7 @@ def to_contiguous_window_dataset( self, window_size=window_size, columns=columns, + anchor_columns=anchor_columns, group_key=group_key, order_key=order_key, stride=stride, diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index d47f2c79c895..ed9da1f7a922 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -44,8 +44,11 @@ class ContiguousWindowDataset(Dataset): * ``pad`` repeats final values and marks repeats in ``is_pad``; * ``error`` rejects the dataset. - ``column_transforms`` convert individual padded column lists and - ``adapter`` can adapt the complete mapping to a model-specific contract. + ``anchor_columns`` limits selected columns to the first row of each window, + which avoids loading repeated context such as observation images. + ``column_transforms`` convert individual column lists and ``adapter`` can + adapt the complete mapping to a model-specific contract. + ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. """ _TAIL_POLICIES = ("drop", "pad", "error") @@ -56,6 +59,7 @@ def __init__( *, window_size, columns=None, + anchor_columns=None, group_key="episode_id", order_key="step_idx", stride=1, @@ -83,6 +87,11 @@ def __init__( raise ValueError("group_key and order_key must not be is_pad.") self.columns = _columns( query, columns, self.group_key, self.order_key) + self.anchor_columns = _anchor_columns(anchor_columns, self.columns) + anchor_column_set = set(self.anchor_columns) + self._window_columns = [ + name for name in self.columns if name not in anchor_column_set + ] self.column_transforms = _column_transforms( column_transforms, self.columns) self.pad_values = _pad_values(pad_values, self.columns) @@ -96,10 +105,6 @@ def __init__( raise ValueError( "ContiguousWindowDataset requires row-tracking.enabled=true.") - self._blob_columns = [ - field.name for field in query._table.fields - if field.name in self.columns and is_blob_type(field.type) - ] index, snapshot_id = _read_window_index( query, self.group_key, self.order_key) self.snapshot_id = snapshot_id @@ -116,7 +121,12 @@ def __len__(self): def __getitem__(self, index): anchor, row_ids = self._resolve_window(index) - return self._sample(anchor, self._read_rows(row_ids)) + rows = self._read_window_rows(row_ids) + anchor_row = ( + self._read_rows(row_ids[:1], self.anchor_columns)[0] + if self.anchor_columns else None + ) + return self._sample(anchor, rows, anchor_row) def __getitems__(self, indices): windows = [self._resolve_window(index) for index in indices] @@ -126,11 +136,22 @@ def __getitems__(self, indices): row_id for _, window_row_ids in windows for row_id in window_row_ids )) - rows_by_id = dict(zip(row_ids, self._read_rows(row_ids))) + rows_by_id = dict(zip(row_ids, self._read_window_rows(row_ids))) + anchor_row_ids = list(dict.fromkeys( + window_row_ids[0] for _, window_row_ids in windows + )) + anchor_rows_by_id = ( + dict(zip( + anchor_row_ids, + self._read_rows(anchor_row_ids, self.anchor_columns), + )) + if self.anchor_columns else {} + ) return [ self._sample( anchor, [rows_by_id[row_id] for row_id in window_row_ids], + anchor_rows_by_id.get(window_row_ids[0]), ) for anchor, window_row_ids in windows ] @@ -147,7 +168,7 @@ def _resolve_window(self, index): row_ids = self._groups[group_index][2] return anchor, row_ids[start:start + valid_count] - def _sample(self, anchor, rows): + def _sample(self, anchor, rows, anchor_row=None): group_index, start, valid_count = anchor group_key, order_values, _ = self._groups[group_index] padding_count = self.window_size - valid_count @@ -160,8 +181,11 @@ def _sample(self, anchor, rows): "is_pad": padding_mask, } for name in self.columns: - values = [row[name] for row in rows] - if padding_count: + if name in self.anchor_columns: + values = [anchor_row[name]] + else: + values = [row[name] for row in rows] + if padding_count and name not in self.anchor_columns: pad_value = self.pad_values.get(name, values[-1]) values.extend( copy.deepcopy(pad_value) for _ in range(padding_count)) @@ -234,7 +258,13 @@ def _build_index(self, index): anchors.append((group_index, start, valid_count)) return groups, anchors - def _read_rows(self, row_ids): + def _read_window_rows(self, row_ids): + if not self._window_columns: + return [{} for _ in row_ids] + return self._read_rows(row_ids, self._window_columns) + + def _read_rows(self, row_ids, columns=None): + columns = self.columns if columns is None else columns query = ScanQuery(self._table) predicate_builder = ( self._table.new_read_builder() @@ -245,14 +275,18 @@ def _read_rows(self, row_ids): ) query._predicate = predicate_builder.is_in( SpecialFields.ROW_ID.name, row_ids) - query._projection = list(self.columns) + query._projection = list(columns) query._include_row_id = True - if self._blob_columns: + blob_columns = [ + field.name for field in self._table.fields + if field.name in columns and is_blob_type(field.type) + ] + if blob_columns: scalar, blobs = query.read_blobs( - self._blob_columns, parallelism=self.blob_parallelism) + blob_columns, parallelism=self.blob_parallelism) rows = scalar.to_pylist() - for name in self._blob_columns: + for name in blob_columns: values = blobs[name] if len(values) != len(rows): raise RuntimeError( @@ -277,9 +311,9 @@ def _read_rows(self, row_ids): return [by_row_id[row_id] for row_id in row_ids] -def _read_window_index(query, group_by, order_by): +def _read_window_index(query, group_key, order_key): index_query = copy.copy(query) - index_query._projection = [group_by, order_by] + index_query._projection = [group_key, order_key] index_query._include_row_id = True read_builder = index_query._configured_read_builder() plan = read_builder.new_scan().plan() @@ -335,6 +369,29 @@ def _columns(query, columns, group_key, order_key): return columns +def _anchor_columns(value, columns): + if value is None: + return [] + if isinstance(value, str): + value = [value] + else: + try: + value = list(value) + except TypeError: + raise TypeError( + "anchor_columns must be a sequence of projected column names.") + if any(not isinstance(name, str) or not name for name in value): + raise TypeError( + "anchor_columns must contain only non-empty column names.") + if len(set(value)) != len(value): + raise ValueError("anchor_columns must not contain duplicates.") + invalid = [name for name in value if name not in columns] + if invalid: + raise ValueError( + "anchor_columns must be included in columns: %s." % invalid) + return value + + def _column_transforms(value, columns): transforms = _mapping(value, "column_transforms") _validate_mapping_columns(transforms, columns, "column_transforms") diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index dc7112a6800b..0dc89e8d81fc 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -145,6 +145,19 @@ def test_reads_blob_payloads_only_when_a_window_is_requested(self): sample["payload"], ) + def test_anchor_columns_read_only_the_window_anchor(self): + table = self._table() + original = ScanQuery._fetch_bodies + with patch.object(ScanQuery, "_fetch_bodies", side_effect=original) as fetch: + dataset = self._dataset(table, anchor_columns=["payload"]) + + sample = dataset[0] + + self.assertEqual([100, 101, 102], sample["value"]) + self.assertEqual([b"episode-b-0"], sample["payload"]) + self.assertEqual(1, fetch.call_count) + self.assertEqual(1, len(fetch.call_args.args[1]["payload"])) + def test_plural_access_coalesces_overlapping_window_reads(self): dataset = self._dataset(self._table()) expected = [dataset[0], dataset[1]] diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py index 40e302f9765a..ca49d3068ef5 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -28,6 +28,7 @@ from pypaimon.benchmark.paired_act import ( IMAGE_COLUMNS, BenchmarkConfig, + _git_head, _paimon_datasets, _shared_normalization, _snapshot_id, @@ -292,6 +293,10 @@ def test_paimon_windows_are_lazy_and_snapshot_pinned(paired_input): assert isinstance(validation, ContiguousWindowDataset) sample_before_append = train[0] assert fetch.call_count == 1 + assert { + name: len(fetch.call_args.args[1][name]) + for name in IMAGE_COLUMNS + } == {name: 1 for name in IMAGE_COLUMNS} scalar, blobs = frames.scan().where( "episode_id = 'train-a' AND frame_index = 5" @@ -341,3 +346,10 @@ def test_tensor_parity_rejects_different_hdf5_bytes( def test_requires_at_least_three_alternating_rounds(): with pytest.raises(ValueError, match="rounds must be at least 3"): BenchmarkConfig(rounds=2) + + +def test_source_commit_falls_back_outside_git_checkout(tmp_path): + with patch( + "pypaimon.benchmark.paired_act.subprocess.check_output", + side_effect=FileNotFoundError): + assert _git_head(tmp_path) == "UNKNOWN"