From 8a850dd87676332b244f108004424a36a69857a5 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 26 Aug 2026 14:51:34 +0800 Subject: [PATCH 01/16] 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: 497/497 AI-Contributed/UT: 314/314 --- 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 b3c59043a8ac..987b3ab42ea7 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -866,6 +866,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 af9189bbae5e..52ee0e52bdce 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -157,7 +157,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 d4491651d8e4..d14de0f84734 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -164,6 +164,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 5988a495dbd1dbba751c2fd480124efdb86fd6b0 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 26 Aug 2026 14:52:44 +0800 Subject: [PATCH 02/16] 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 12e24632bc3e..984ef38ba932 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -255,6 +255,10 @@ def read_requirements(): 'torch': [ 'torch', ], + 'act': [ + 'lerobot==0.4.4', + 'Pillow', + ], 'daft': [ 'daft>=0.7.6; python_version>="3.10"', ], From 94d7d0dcb74f6d2cd158a57926151789ff45ee81 Mon Sep 17 00:00:00 2001 From: Yann Date: Sat, 29 Aug 2026 22:43:07 +0800 Subject: [PATCH 03/16] 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 52ee0e52bdce..f6e7c87617e1 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -173,6 +173,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", @@ -185,7 +186,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 d14de0f84734..00bcb0f4fc31 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -169,6 +169,7 @@ def to_contiguous_window_dataset( *, window_size, columns=None, + anchor_columns=None, group_key="episode_id", order_key="step_idx", stride=1, @@ -180,8 +181,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. """ @@ -194,6 +196,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" From 2af914fe6dc6992af9aa8d1dc360ef32334a9216 Mon Sep 17 00:00:00 2001 From: Yann Date: Sun, 30 Aug 2026 01:14:52 +0800 Subject: [PATCH 04/16] perf(python): coalesce ACT physical fetches Fetch several physical batches at once while preserving the logical batch boundaries consumed by ACT training. Keep warmup isolation, cursor replay, and snapshot semantics unchanged. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 102/102 AI-Contributed/UT: 105/105 --- docs/docs/pypaimon/robomind-act-benchmark.md | 8 ++ .../pypaimon/benchmark/act_harness.py | 90 +++++++++------ .../pypaimon/benchmark/paired_act.py | 4 + .../tests/paired_act_benchmark_test.py | 105 ++++++++++++++++++ 4 files changed, 176 insertions(+), 31 deletions(-) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index 664465fc111e..4fd06f70e0ec 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -59,6 +59,14 @@ 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. +The logical training batch remains unchanged, while `fetch_batches` controls +how many consecutive logical batches are requested from the Dataset together. +The default of four lets Paimon coalesce eight samples when the logical batch +size is two, then yields the original four ordered batches. The transient +buffer is discarded after that physical fetch; checkpoint recovery resumes +from the next logical-batch cursor and reconstructs it. Larger values trade +fewer Paimon reads for higher per-process 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 diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act_harness.py index dd874406579c..65338f2e903b 100644 --- a/paimon-python/pypaimon/benchmark/act_harness.py +++ b/paimon-python/pypaimon/benchmark/act_harness.py @@ -30,7 +30,7 @@ import torch import torch.nn.functional as functional from PIL import Image -from torch.utils.data import DataLoader, Dataset +from torch.utils.data import Dataset, default_collate CAMERA_KEYS = ( @@ -54,6 +54,7 @@ class BenchmarkConfig: weight_decay: float = 1e-4 warmup_batches: int = 1 loader_batches: int = 4 + fetch_batches: int = 4 rounds: int = 3 def __post_init__(self): @@ -65,6 +66,7 @@ def __post_init__(self): "image_width", "warmup_batches", "loader_batches", + "fetch_batches", ) for name in positive_ints: value = getattr(self, name) @@ -288,25 +290,30 @@ def run_backend( 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, + warmup_sample_count = config.warmup_batches * config.batch_size + warmup_iterator = _iter_logical_batches( + train_dataset, + plan.loader_indices[:warmup_sample_count], + logical_batch_size=config.batch_size, + fetch_batches=1, ) - iterator = iter(loader) first_batch_started = time.monotonic() - first_batch = next(iterator) + first_batch = next(warmup_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) + validate_act_batch(next(warmup_iterator), config) + loader_iterator = _iter_logical_batches( + train_dataset, + plan.loader_indices[warmup_sample_count:], + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + ) loader_started = time.monotonic() loader_sample_count = 0 for _ in range(config.loader_batches): - batch = next(iterator) + batch = next(loader_iterator) validate_act_batch(batch, config) loader_sample_count += len(batch["sample_id"]) loader_seconds = time.monotonic() - loader_started @@ -322,15 +329,14 @@ def run_backend( 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): + for step, batch in enumerate(_iter_logical_batches( + train_dataset, + plan.train_indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + ), 1): step_started = time.monotonic() model_batch = build_lerobot_batch(batch, config) optimizer.zero_grad(set_to_none=True) @@ -360,12 +366,12 @@ def run_backend( # 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, - ))) + validation_batch = next(_iter_logical_batches( + validation_dataset, + plan.validation_indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + )) with torch.no_grad(): validation_loss, _ = policy(build_lerobot_batch( validation_batch, config)) @@ -405,13 +411,15 @@ def _measure_python_peak(dataset_factory, plan, config): 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, - ))) + indices = plan.loader_indices[ + :config.batch_size * config.fetch_batches + ] + next(_iter_logical_batches( + train_dataset, + indices, + logical_batch_size=config.batch_size, + fetch_batches=config.fetch_batches, + )) _, peak = tracemalloc.get_traced_memory() return peak finally: @@ -468,3 +476,23 @@ def __getitems__(self, indices): if getitems is not None: return getitems(source_indices) return [self._dataset[index] for index in source_indices] + + +def _iter_logical_batches( + dataset, indices, *, logical_batch_size, fetch_batches): + logical_batch_size = _positive_int( + logical_batch_size, "logical_batch_size") + fetch_batches = _positive_int(fetch_batches, "fetch_batches") + if len(indices) % logical_batch_size: + raise ValueError("indices must contain complete logical batches.") + physical_size = logical_batch_size * fetch_batches + getitems = getattr(dataset, "__getitems__", None) + for offset in range(0, len(indices), physical_size): + physical_indices = list(indices[offset:offset + physical_size]) + if getitems is None: + samples = [dataset[index] for index in physical_indices] + else: + samples = getitems(physical_indices) + for logical_offset in range(0, len(samples), logical_batch_size): + yield default_collate( + samples[logical_offset:logical_offset + logical_batch_size]) diff --git a/paimon-python/pypaimon/benchmark/paired_act.py b/paimon-python/pypaimon/benchmark/paired_act.py index 93cfa512d05b..b237d8b63708 100644 --- a/paimon-python/pypaimon/benchmark/paired_act.py +++ b/paimon-python/pypaimon/benchmark/paired_act.py @@ -826,6 +826,9 @@ def main(argv=None): parser.add_argument( "--loader-batches", type=int, default=BenchmarkConfig.loader_batches, help="Batches used for DataLoader throughput measurement.") + parser.add_argument( + "--fetch-batches", type=int, default=BenchmarkConfig.fetch_batches, + help="Logical batches coalesced into one physical dataset fetch.") parser.add_argument( "--rounds", type=int, default=BenchmarkConfig.rounds, help="Alternating backend rounds; must be at least three.") @@ -841,6 +844,7 @@ def main(argv=None): weight_decay=args.weight_decay, warmup_batches=args.warmup_batches, loader_batches=args.loader_batches, + fetch_batches=args.fetch_batches, rounds=args.rounds, ) report = run( diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py index ca49d3068ef5..c990e6d3d1bb 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -25,6 +25,8 @@ from PIL import Image import pypaimon.multimodal as pmm +import pypaimon.benchmark.act_harness as act_harness +import pypaimon.benchmark.paired_act as paired_act from pypaimon.benchmark.paired_act import ( IMAGE_COLUMNS, BenchmarkConfig, @@ -57,6 +59,44 @@ def __getitems__(self, indices): assert dataset.__getitems__([0, 2]) == ["sample-7", "sample-5"] +def test_logical_batches_coalesce_one_physical_fetch(): + class BatchDataset: + def __init__(self): + self.calls = [] + + def __getitems__(self, indices): + self.calls.append(list(indices)) + return [{"value": torch.tensor(index)} for index in indices] + + dataset = BatchDataset() + + batches = list(act_harness._iter_logical_batches( + dataset, + tuple(range(8)), + logical_batch_size=2, + fetch_batches=4, + )) + + assert dataset.calls == [list(range(8))] + assert [batch["value"].tolist() for batch in batches] == [ + [0, 1], [2, 3], [4, 5], [6, 7], + ] + + +def test_logical_batches_reject_partial_checkpoint_tail(): + class BatchDataset: + def __getitems__(self, indices): + return [{"value": torch.tensor(index)} for index in indices] + + with pytest.raises(ValueError, match="complete logical batches"): + list(act_harness._iter_logical_batches( + BatchDataset(), + tuple(range(9)), + logical_batch_size=2, + fetch_batches=4, + )) + + def test_backend_times_without_tracemalloc_and_measures_memory_separately(): states = [] config = BenchmarkConfig( @@ -105,6 +145,58 @@ def __getitem__(self, index): "python-tracemalloc-separate-dataset-first-batch") +def test_backend_coalesces_timed_loader_fetches(): + config = BenchmarkConfig( + seed=11, + action_horizon=1, + batch_size=2, + optimizer_steps=1, + image_height=2, + image_width=2, + warmup_batches=1, + loader_batches=4, + fetch_batches=4, + rounds=3, + ) + + class BatchDataset(torch.utils.data.Dataset): + def __init__(self): + self.calls = [] + + def __len__(self): + return 16 + + def __getitem__(self, index): + 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), + } + + def __getitems__(self, indices): + self.calls.append(list(indices)) + return [self[index] for index in indices] + + dataset = BatchDataset() + plan = build_window_plan(len(dataset), len(dataset), config) + + run_backend( + "test", + 1, + lambda: (dataset, dataset), + plan, + config, + "sequence-sha256", + policy_factory=_policy_factory, + ) + + assert any(len(indices) == 8 for indices in dataset.calls) + + def _jpeg(value): buffer = BytesIO() Image.fromarray(np.full((8, 10, 3), value, dtype=np.uint8)).save( @@ -348,6 +440,19 @@ def test_requires_at_least_three_alternating_rounds(): BenchmarkConfig(rounds=2) +def test_fetch_batches_must_be_positive(): + assert BenchmarkConfig(fetch_batches=4).fetch_batches == 4 + with pytest.raises(ValueError, match="fetch_batches must be a positive int"): + BenchmarkConfig(fetch_batches=0) + + +def test_cli_documents_physical_fetch_batches(capsys): + with pytest.raises(SystemExit): + paired_act.main(["--help"]) + + assert "--fetch-batches" in capsys.readouterr().out + + def test_source_commit_falls_back_outside_git_checkout(tmp_path): with patch( "pypaimon.benchmark.paired_act.subprocess.check_output", From 5aef6d0897195a80b2b6bd334e6a489399692c3f Mon Sep 17 00:00:00 2001 From: Yann Date: Sun, 30 Aug 2026 01:23:32 +0800 Subject: [PATCH 05/16] perf(python): tune ACT physical fetch depth Use eight logical batches per physical fetch, based on the current Vortex-backed RoboMIND layout. Preserve logical batch boundaries while documenting the per-worker memory tradeoff. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 6/6 AI-Contributed/UT: 4/4 --- docs/docs/pypaimon/robomind-act-benchmark.md | 4 ++-- paimon-python/pypaimon/benchmark/act_harness.py | 2 +- paimon-python/pypaimon/tests/paired_act_benchmark_test.py | 4 ++++ 3 files changed, 7 insertions(+), 3 deletions(-) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index 4fd06f70e0ec..e53e478f4aec 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -61,8 +61,8 @@ materializing episodes in memory. The logical training batch remains unchanged, while `fetch_batches` controls how many consecutive logical batches are requested from the Dataset together. -The default of four lets Paimon coalesce eight samples when the logical batch -size is two, then yields the original four ordered batches. The transient +The default of eight lets Paimon coalesce 16 samples when the logical batch +size is two, then yields the original eight ordered batches. The transient buffer is discarded after that physical fetch; checkpoint recovery resumes from the next logical-batch cursor and reconstructs it. Larger values trade fewer Paimon reads for higher per-process memory. diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act_harness.py index 65338f2e903b..ddcc5f37e82c 100644 --- a/paimon-python/pypaimon/benchmark/act_harness.py +++ b/paimon-python/pypaimon/benchmark/act_harness.py @@ -54,7 +54,7 @@ class BenchmarkConfig: weight_decay: float = 1e-4 warmup_batches: int = 1 loader_batches: int = 4 - fetch_batches: int = 4 + fetch_batches: int = 8 rounds: int = 3 def __post_init__(self): diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py index c990e6d3d1bb..7ae8e90377aa 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -49,6 +49,10 @@ h5py = pytest.importorskip("h5py") +def test_default_fetch_group_covers_eight_logical_batches(): + assert BenchmarkConfig().fetch_batches == 8 + + def test_sequence_dataset_forwards_plural_access(): class BatchDataset: def __getitems__(self, indices): From 6a5fc461f83df4c298d48679cfad7bf0ea1359c3 Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 1 Sep 2026 11:20:12 +0800 Subject: [PATCH 06/16] fix(python): align ACT benchmark with latest master Reuse the guarded LeRobot dependency set for the ACT extra, clarify batch-fetch measurement terminology, and strengthen the timed plural-fetch assertion. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 136/136 AI-Contributed/UT: 54/54 --- docs/docs/pypaimon/robomind-act-benchmark.md | 19 +++--- .../pypaimon/benchmark/act_harness.py | 66 +++++++------------ .../pypaimon/benchmark/paired_act.py | 28 ++++---- .../tests/paired_act_benchmark_test.py | 54 +++++++-------- paimon-python/setup.py | 23 +++---- 5 files changed, 81 insertions(+), 109 deletions(-) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index e53e478f4aec..1e1faad4a07a 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -44,7 +44,7 @@ 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 +A seeded window plan fixes every warmup, batch-fetch, training, and validation 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. @@ -53,19 +53,19 @@ 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 -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 +`ContiguousWindowDataset.__getitems__` coalesces overlapping row IDs into one +payload read. The 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. The logical training batch remains unchanged, while `fetch_batches` controls how many consecutive logical batches are requested from the Dataset together. The default of eight lets Paimon coalesce 16 samples when the logical batch size is two, then yields the original eight ordered batches. The transient -buffer is discarded after that physical fetch; checkpoint recovery resumes -from the next logical-batch cursor and reconstructs it. Larger values trade -fewer Paimon reads for higher per-process memory. +buffer is discarded after that physical fetch. Recovery behavior is not covered +by this benchmark. Larger values trade fewer Paimon reads for higher per-process +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 @@ -78,7 +78,8 @@ 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 +- first-batch latency, dataset batch-fetch 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; diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act_harness.py index ddcc5f37e82c..8b121ec5cd74 100644 --- a/paimon-python/pypaimon/benchmark/act_harness.py +++ b/paimon-python/pypaimon/benchmark/act_harness.py @@ -30,7 +30,7 @@ import torch import torch.nn.functional as functional from PIL import Image -from torch.utils.data import Dataset, default_collate +from torch.utils.data import default_collate CAMERA_KEYS = ( @@ -53,7 +53,7 @@ class BenchmarkConfig: learning_rate: float = 1e-4 weight_decay: float = 1e-4 warmup_batches: int = 1 - loader_batches: int = 4 + timed_batches: int = 4 fetch_batches: int = 8 rounds: int = 3 @@ -65,7 +65,7 @@ def __post_init__(self): "image_height", "image_width", "warmup_batches", - "loader_batches", + "timed_batches", "fetch_batches", ) for name in positive_ints: @@ -95,7 +95,7 @@ class WindowPlan: """Explicit window indices consumed identically by both backends.""" seed: int - loader_indices: tuple + measurement_indices: tuple train_indices: tuple validation_indices: tuple @@ -108,25 +108,25 @@ def sha256(self): def to_dict(self): return { "seed": self.seed, - "loader_indices": list(self.loader_indices), + "measurement_indices": list(self.measurement_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.""" + """Build stable indices for training, validation, and batch-fetch 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 + batch_fetch_count = ( + config.warmup_batches + config.timed_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)), + measurement_indices=tuple(_repeat_permutations( + train_window_count, batch_fetch_count, config.seed + 1)), train_indices=tuple(_repeat_permutations( train_window_count, train_count, config.seed + 2)), validation_indices=tuple(_repeat_permutations( @@ -293,7 +293,7 @@ def run_backend( warmup_sample_count = config.warmup_batches * config.batch_size warmup_iterator = _iter_logical_batches( train_dataset, - plan.loader_indices[:warmup_sample_count], + plan.measurement_indices[:warmup_sample_count], logical_batch_size=config.batch_size, fetch_batches=1, ) @@ -304,19 +304,19 @@ def run_backend( for _ in range(config.warmup_batches - 1): validate_act_batch(next(warmup_iterator), config) - loader_iterator = _iter_logical_batches( + batch_fetch_iterator = _iter_logical_batches( train_dataset, - plan.loader_indices[warmup_sample_count:], + plan.measurement_indices[warmup_sample_count:], logical_batch_size=config.batch_size, fetch_batches=config.fetch_batches, ) - loader_started = time.monotonic() - loader_sample_count = 0 - for _ in range(config.loader_batches): - batch = next(loader_iterator) + batch_fetch_started = time.monotonic() + batch_fetch_sample_count = 0 + for _ in range(config.timed_batches): + batch = next(batch_fetch_iterator) validate_act_batch(batch, config) - loader_sample_count += len(batch["sample_id"]) - loader_seconds = time.monotonic() - loader_started + batch_fetch_sample_count += len(batch["sample_id"]) + batch_fetch_seconds = time.monotonic() - batch_fetch_started _seed_everything(config.seed) policy, model = policy_factory(config) @@ -392,9 +392,10 @@ def run_backend( "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, + "batch_fetch_samples": batch_fetch_sample_count, + "batch_fetch_s": batch_fetch_seconds, + "batch_fetch_samples_per_s": ( + batch_fetch_sample_count / batch_fetch_seconds), "fixed_steps_s": fixed_steps_s, "train_loss": [item["total"] for item in losses], "train_trace": losses, @@ -411,7 +412,7 @@ def _measure_python_peak(dataset_factory, plan, config): tracemalloc.start() try: train_dataset, _ = dataset_factory() - indices = plan.loader_indices[ + indices = plan.measurement_indices[ :config.batch_size * config.fetch_batches ] next(_iter_logical_batches( @@ -459,25 +460,6 @@ def _positive_int(value, 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] - - def _iter_logical_batches( dataset, indices, *, logical_batch_size, fetch_batches): logical_batch_size = _positive_int( diff --git a/paimon-python/pypaimon/benchmark/paired_act.py b/paimon-python/pypaimon/benchmark/paired_act.py index b237d8b63708..b2bb30ebddd6 100644 --- a/paimon-python/pypaimon/benchmark/paired_act.py +++ b/paimon-python/pypaimon/benchmark/paired_act.py @@ -117,7 +117,7 @@ def __getitem__(self, anchor): slice(anchor, anchor + self.action_horizon), ) images = np.stack([ - _decode_hdf5_image(h5[field][anchor]) + _decode_image(h5[field][anchor]) for field in HDF5_IMAGE_FIELDS ]) qpos = ( @@ -320,7 +320,6 @@ def run( "config": config.to_dict(), "cache_control": "uncontrolled", "device": "cpu", - "data_loader_workers": 0, }, "normalization": normalization_metadata, "window_plan": { @@ -356,7 +355,7 @@ def run( "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 " + "GPU, multi-worker dataset loading, 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-" @@ -574,7 +573,7 @@ def _select_episode(source_by_id, split, requested, action_horizon): def _tensor_parity(hdf5_datasets, paimon_datasets, plan): comparisons = ( ("train", hdf5_datasets[0], paimon_datasets[0], - sorted(set(plan.loader_indices + plan.train_indices))), + sorted(set(plan.measurement_indices + plan.train_indices))), ("validation", hdf5_datasets[1], paimon_datasets[1], sorted(set(plan.validation_indices))), ) @@ -651,7 +650,7 @@ def _summarize(runs): metrics = ( "dataset_build_s", "first_batch_s", - "dataloader_samples_per_s", + "batch_fetch_samples_per_s", "fixed_steps_s", "validation_loss", "python_peak_allocated_bytes", @@ -670,9 +669,9 @@ def _summarize(runs): def _sample_sequence_sha256(train_episode_id, validation_episode_id, plan): value = { - "loader": [ + "batch_fetch": [ "%s#%d" % (train_episode_id, index) - for index in plan.loader_indices + for index in plan.measurement_indices ], "train": [ "%s#%d" % (train_episode_id, index) @@ -698,10 +697,6 @@ def _read_vectors(h5, fields, selection, dtype=np.float32): return value -def _decode_hdf5_image(value): - return _decode_image(value) - - def _decode_image(value): payload = ( bytes(value) @@ -804,7 +799,7 @@ def main(argv=None): help="Number of contiguous action rows in each sample.") parser.add_argument( "--batch-size", type=int, default=BenchmarkConfig.batch_size, - help="Shared DataLoader batch size.") + help="Shared logical batch size.") parser.add_argument( "--optimizer-steps", type=int, default=BenchmarkConfig.optimizer_steps, help="Fixed optimizer steps per backend run.") @@ -822,10 +817,11 @@ def main(argv=None): help="Shared AdamW weight decay.") parser.add_argument( "--warmup-batches", type=int, default=BenchmarkConfig.warmup_batches, - help="DataLoader batches consumed before timing.") + help="Logical batches consumed before batch-fetch timing.") parser.add_argument( - "--loader-batches", type=int, default=BenchmarkConfig.loader_batches, - help="Batches used for DataLoader throughput measurement.") + "--timed-batches", type=int, + default=BenchmarkConfig.timed_batches, + help="Logical batches used for dataset batch-fetch throughput.") parser.add_argument( "--fetch-batches", type=int, default=BenchmarkConfig.fetch_batches, help="Logical batches coalesced into one physical dataset fetch.") @@ -843,7 +839,7 @@ def main(argv=None): learning_rate=args.learning_rate, weight_decay=args.weight_decay, warmup_batches=args.warmup_batches, - loader_batches=args.loader_batches, + timed_batches=args.timed_batches, fetch_batches=args.fetch_batches, rounds=args.rounds, ) diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py index 7ae8e90377aa..5c87d48eb8a3 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -36,11 +36,7 @@ _snapshot_id, run, ) -from pypaimon.benchmark.act_harness import ( - _SequenceDataset, - build_window_plan, - run_backend, -) +from pypaimon.benchmark.act_harness import 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 @@ -53,16 +49,6 @@ def test_default_fetch_group_covers_eight_logical_batches(): assert BenchmarkConfig().fetch_batches == 8 -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_logical_batches_coalesce_one_physical_fetch(): class BatchDataset: def __init__(self): @@ -87,7 +73,7 @@ def __getitems__(self, indices): ] -def test_logical_batches_reject_partial_checkpoint_tail(): +def test_logical_batches_reject_incomplete_batch_tail(): class BatchDataset: def __getitems__(self, indices): return [{"value": torch.tensor(index)} for index in indices] @@ -111,7 +97,7 @@ def test_backend_times_without_tracemalloc_and_measures_memory_separately(): image_height=2, image_width=2, warmup_batches=1, - loader_batches=1, + timed_batches=1, rounds=3, ) @@ -149,7 +135,7 @@ def __getitem__(self, index): "python-tracemalloc-separate-dataset-first-batch") -def test_backend_coalesces_timed_loader_fetches(): +def test_backend_coalesces_timed_batch_fetches(): config = BenchmarkConfig( seed=11, action_horizon=1, @@ -158,7 +144,7 @@ def test_backend_coalesces_timed_loader_fetches(): image_height=2, image_width=2, warmup_batches=1, - loader_batches=4, + timed_batches=4, fetch_batches=4, rounds=3, ) @@ -188,17 +174,23 @@ def __getitems__(self, indices): dataset = BatchDataset() plan = build_window_plan(len(dataset), len(dataset), config) - run_backend( - "test", - 1, - lambda: (dataset, dataset), - plan, - config, - "sequence-sha256", - policy_factory=_policy_factory, - ) + with patch.object(act_harness, "_measure_python_peak", return_value=0): + run_backend( + "test", + 1, + lambda: (dataset, dataset), + plan, + config, + "sequence-sha256", + policy_factory=_policy_factory, + ) - assert any(len(indices) == 8 for indices in dataset.calls) + assert dataset.calls == [ + list(plan.measurement_indices[:2]), + list(plan.measurement_indices[2:10]), + list(plan.train_indices), + list(plan.validation_indices), + ] def _jpeg(value): @@ -279,7 +271,7 @@ def test_runs_three_alternating_rounds_with_one_shared_contract( image_height=8, image_width=10, warmup_batches=1, - loader_batches=2, + timed_batches=2, rounds=3, ) @@ -333,7 +325,7 @@ def test_runs_three_alternating_rounds_with_one_shared_contract( assert report["summary"][backend]["round_count"] == 3 for metric in ( "first_batch_s", - "dataloader_samples_per_s", + "batch_fetch_samples_per_s", "fixed_steps_s", "python_peak_allocated_bytes"): assert set(report["summary"][backend][metric]) == { diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 984ef38ba932..382371733ab7 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -216,6 +216,15 @@ def read_requirements(): install_requires = read_requirements() +LEROBOT_DEPENDENCIES = [ + # datasets 4.1+ may select PyArrow 21+, while PyPaimon currently + # supports PyArrow <20. Pandas 2.2.2+ supports NumPy 2.x selected + # by LeRobot's media dependencies. + 'datasets>=4,<4.1; python_version>="3.10"', + 'pandas>=2.2.2,<3; python_version>="3.10"', + 'lerobot>=0.4.4,<0.5; python_version>="3.10"', +] + long_description = "See Apache Paimon Python API \ [Doc](https://paimon.apache.org/docs/master/pypaimon/python-api/) for usage." @@ -241,23 +250,15 @@ def read_requirements(): # rosbags is pure Python and does not require a ROS installation. 'rosbags>=0.11.5,<0.12; python_version>="3.10"', ], - 'lerobot': [ - # datasets 4.1+ may select PyArrow 21+, while PyPaimon currently - # supports PyArrow <20. Pandas 2.2.2+ supports NumPy 2.x selected - # by LeRobot's media dependencies. - 'datasets>=4,<4.1; python_version>="3.10"', - 'pandas>=2.2.2,<3; python_version>="3.10"', - 'lerobot>=0.4.4,<0.5; python_version>="3.10"', - ], + 'lerobot': LEROBOT_DEPENDENCIES, 'ray': [ 'ray>=2.10,<3; python_version>="3.8"', ], 'torch': [ 'torch', ], - 'act': [ - 'lerobot==0.4.4', - 'Pillow', + 'act': LEROBOT_DEPENDENCIES + [ + 'Pillow; python_version>="3.10"', ], 'daft': [ 'daft>=0.7.6; python_version>="3.10"', From 6c8b2de9f266f396e150375fba446121522f90ef Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 1 Sep 2026 14:34:28 +0800 Subject: [PATCH 07/16] fix(python): isolate ACT benchmark test dependencies Skip the optional ACT test module before importing unavailable dependencies and keep its synthetic RoboMIND fixture independent of Vortex. Co-Authored-By: Codex AI-Model: gpt-5 AI-Contributed/Feature: 0/0 AI-Contributed/UT: 23/23 --- .../tests/paired_act_benchmark_test.py | 23 +++++++++++++------ 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py index 5c87d48eb8a3..6e4e96002003 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -14,6 +14,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +# ruff: noqa: E402 + import json import tracemalloc from io import BytesIO @@ -21,8 +23,11 @@ import numpy as np import pytest -import torch -from PIL import Image + + +torch = pytest.importorskip("torch") +Image = pytest.importorskip("PIL.Image") +h5py = pytest.importorskip("h5py") import pypaimon.multimodal as pmm import pypaimon.benchmark.act_harness as act_harness @@ -42,9 +47,6 @@ from pypaimon.sample import robomind_agilex as agilex -h5py = pytest.importorskip("h5py") - - def test_default_fetch_group_covers_eight_logical_batches(): assert BenchmarkConfig().fetch_batches == 8 @@ -223,12 +225,16 @@ def _write_episode(root, split, name, offset, frames=6): @pytest.fixture -def paired_input(tmp_path): +def paired_input(tmp_path, monkeypatch): 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" + monkeypatch.setattr(agilex, "TABLE_OPTIONS", { + **agilex.TABLE_OPTIONS, + "vector.file.format": "parquet", + }) agilex.ingest_local(root, warehouse, batch_size=2) agilex.backfill_canonical_action( warehouse, statistics_version="paired-test@1") @@ -343,13 +349,16 @@ def test_runs_three_alternating_rounds_with_one_shared_contract( by_backend["paimon"]["validation_loss"]) -def test_paimon_windows_are_lazy_and_snapshot_pinned(paired_input): +def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( + 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) + assert frames.raw_table.table_schema.options["vector.file.format"] == ( + "parquet") snapshot_id = _snapshot_id(frames) normalization, _ = _shared_normalization( agilex.discover_episodes(input_root), From 1f53ff310b82db692b912e8dc8a27af9aec3f865 Mon Sep 17 00:00:00 2001 From: Yann Date: Tue, 1 Sep 2026 20:06:51 +0800 Subject: [PATCH 08/16] refactor(python): modularize ACT benchmark execution Separate experiment preparation, backend execution, and result comparison while preserving a shared ACT workload contract. Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 3667/3667 AI-Contributed/UT: 1330/1330 --- docs/docs/pypaimon/robomind-act-benchmark.md | 211 +++-- .../pypaimon/benchmark/act/__init__.py | 17 + .../pypaimon/benchmark/act/__main__.py | 187 ++++ .../pypaimon/benchmark/act/compare.py | 231 +++++ .../benchmark/act/default_experiment.json | 22 + .../pypaimon/benchmark/act/experiment.py | 42 + .../{act_harness.py => act/harness.py} | 113 ++- paimon-python/pypaimon/benchmark/act/hdf5.py | 196 ++++ .../pypaimon/benchmark/act/paimon.py | 145 +++ .../pypaimon/benchmark/act/runner.py | 691 ++++++++++++++ .../pypaimon/benchmark/paired_act.py | 865 ------------------ .../pypaimon/tests/act_benchmark_test.py | 233 +++++ ...t_benchmark_test.py => act_runner_test.py} | 411 ++++++--- paimon-python/setup.py | 7 +- 14 files changed, 2312 insertions(+), 1059 deletions(-) create mode 100644 paimon-python/pypaimon/benchmark/act/__init__.py create mode 100644 paimon-python/pypaimon/benchmark/act/__main__.py create mode 100644 paimon-python/pypaimon/benchmark/act/compare.py create mode 100644 paimon-python/pypaimon/benchmark/act/default_experiment.json create mode 100644 paimon-python/pypaimon/benchmark/act/experiment.py rename paimon-python/pypaimon/benchmark/{act_harness.py => act/harness.py} (77%) create mode 100644 paimon-python/pypaimon/benchmark/act/hdf5.py create mode 100644 paimon-python/pypaimon/benchmark/act/paimon.py create mode 100644 paimon-python/pypaimon/benchmark/act/runner.py delete mode 100644 paimon-python/pypaimon/benchmark/paired_act.py create mode 100644 paimon-python/pypaimon/tests/act_benchmark_test.py rename paimon-python/pypaimon/tests/{paired_act_benchmark_test.py => act_runner_test.py} (50%) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index 1e1faad4a07a..bb40a8c4277f 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -1,5 +1,5 @@ --- -title: "RoboMIND Paired ACT Benchmark" +title: "RoboMIND ACT Storage Benchmark" sidebar_position: 8 --- @@ -22,71 +22,160 @@ specific language governing permissions and limitations under the License. --> -# RoboMIND Paired ACT Benchmark +# RoboMIND ACT Storage 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: +This benchmark measures the same CPU LeRobot ACT training workload over an +original RoboMIND AgileX HDF5 dataset or an already ingested and +canonical-action-backfilled Paimon warehouse. Ingestion and backfill are outside +the timed scope. + +The backends run independently. A resolved experiment document preserves the +shared configuration, normalization, seed, episode selection, Paimon snapshot, +and logical window sequence. Result comparison verifies that contract before it +calculates performance ratios. + +## Install ```shell pip install 'pypaimon[act,hdf5]' -python -m pypaimon.benchmark.paired_act \ +``` + +## 1. Prepare the experiment + +```shell +python -m pypaimon.benchmark.act prepare \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --warehouse /data/warehouse \ + --output /data/results/experiment.json +``` + +Preparation is not timed. It verifies that HDF5 discovery matches the Paimon +episodes table, checks versioned action statistics against train-only HDF5 +moments, selects eligible train and validation episodes, pins the frames +snapshot, and materializes deterministic measurement, training, and validation +window indices. + +Without `--experiment`, preparation starts from the packaged +`default_experiment.json`. A custom JSON file can change the defaults, and +individual values can be overridden on the command line: + +```shell +python -m pypaimon.benchmark.act prepare \ + --experiment my-experiment.json \ --input /data/RoboMIND/h5_agilex_3rgb \ --warehouse /data/warehouse \ - --report /data/results/paired-act.json + --action-horizon 32 \ + --batch-size 2 \ + --fetch-batches 8 \ + --rounds 3 \ + --output /data/results/experiment.json +``` + +The resolved experiment embeds the effective parameters as well as: + +- portable source episode metadata and its SHA-256; +- normalization values, scope, version, frame count, and SHA-256; +- selected train and validation episode IDs; +- every logical window index, the window-plan SHA-256, and the + episode-qualified sample-sequence SHA-256; +- the Paimon database, frames table, and pinned snapshot ID. + +## 2. Run each backend + +```shell +python -m pypaimon.benchmark.act run \ + --backend hdf5 \ + --experiment /data/results/experiment.json \ + --input /data/RoboMIND/h5_agilex_3rgb \ + --results-dir /data/results + +python -m pypaimon.benchmark.act run \ + --backend paimon \ + --experiment /data/results/experiment.json \ + --warehouse /data/warehouse \ + --results-dir /data/results ``` -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, batch-fetch, training, and validation -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. -`ContiguousWindowDataset.__getitems__` coalesces overlapping row IDs into one -payload read. The 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. - -The logical training batch remains unchanged, while `fetch_batches` controls -how many consecutive logical batches are requested from the Dataset together. -The default of eight lets Paimon coalesce 16 samples when the logical batch -size is two, then yields the original eight ordered batches. The transient -buffer is discarded after that physical fetch. Recovery behavior is not covered -by this benchmark. Larger values trade fewer Paimon reads for higher per-process -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, dataset batch-fetch 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. +Use `--output` to choose an exact result path. Otherwise the command writes an +automatically named JSON file below `--results-dir` and prints its absolute +path as a compact JSON object. + +Each result contains the complete resolved experiment and experiment SHA-256, +backend identity, runtime environment, model metadata, planned-sample tensor +fingerprint, three or more raw measurement rounds, and median/minimum/maximum +summary metrics. + +Both adapters produce the same shared sample contract. State and camera images +come from the anchor frame; action covers the complete horizon. HDF5 reads a +window on demand from one episode file. Paimon uses a lazy, snapshot-pinned +`ContiguousWindowDataset`; image columns are anchor-only, and plural +`__getitems__` access coalesces multiple logical batches into a physical +fetch before splitting them back into the unchanged model batch size. + +## 3. Compare results + +Compare explicit files: + +```shell +python -m pypaimon.benchmark.act compare \ + /data/results/robomind-act-hdf5-20260901T010000Z-a1b2c3d4.json \ + /data/results/robomind-act-paimon-20260901T011000Z-e5f6a7b8.json \ + --output /data/results/comparison.json +``` + +Or discover all ACT result documents in a directory: + +```shell +python -m pypaimon.benchmark.act compare \ + --results-dir /data/results \ + --output /data/results/comparison.json +``` + +Directory discovery ignores experiment and prior comparison JSON files. +Results are grouped by experiment SHA-256. Different experiments remain +separate entries in one comparison artifact; only compatible repeated results +for the same experiment and backend are aggregated. + +Within one experiment group, comparison requires identical runtime environment, +model metadata, tensor fingerprint, train-loss trace, and validation-loss trace. +An environment mismatch marks the group `INCOMPATIBLE`; a model, tensor, or +loss mismatch marks it `FAILED`. Neither case produces performance ratios. + +For compatible HDF5 and Paimon results, higher-is-better metrics report +`paimon_over_hdf5`. Lower-is-better latency, time, and memory metrics report +`hdf5_over_paimon`, which is the Paimon speedup or reduction factor. + +## Measurements + +Every backend repeat records: + +- dataset construction time; +- first-batch latency after construction; +- batch-fetch samples per second after warm-up; +- fixed ACT optimizer-step time and per-step loss trace; +- validation loss; +- total measured wall time; +- Python peak allocation from a separate dataset-first-batch replay. + +The shared harness resets Python, NumPy, and Torch random generators before +model construction and enables deterministic Torch algorithms. The logical +window plan is explicit rather than delegated to a streaming reader. + +Python peak allocation uses `tracemalloc` after wall-clock measurement so +tracing overhead does not distort throughput. It does not include every native +Arrow or Torch allocation. The benchmark does not drop the OS page cache. +GPU, multi-worker loading, distributed training, recovery, and policy quality +remain outside this benchmark. + +## Code organization + +- `benchmark.act.harness`: shared ACT tensors, model, trainer, window plan, and + measurement lifecycle; +- `benchmark.act.hdf5`: HDF5 window dataset and train normalization moments; +- `benchmark.act.paimon`: Paimon adapter, snapshot-pinned datasets, and + versioned statistics access; +- `benchmark.act.runner`: experiment preparation and one-backend execution; +- `benchmark.act.compare`: result discovery, compatibility checks, grouping by + experiment, and aggregation of compatible repeated runs; +- `benchmark.act.__main__`: the `prepare`, `run`, and `compare` + command-line interface. diff --git a/paimon-python/pypaimon/benchmark/act/__init__.py b/paimon-python/pypaimon/benchmark/act/__init__.py new file mode 100644 index 000000000000..f6224db90413 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/__init__.py @@ -0,0 +1,17 @@ +# 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. + +"""ACT training benchmark backends and result comparison.""" diff --git a/paimon-python/pypaimon/benchmark/act/__main__.py b/paimon-python/pypaimon/benchmark/act/__main__.py new file mode 100644 index 000000000000..686708e42072 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/__main__.py @@ -0,0 +1,187 @@ +# 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. + +"""Command-line entry point for ACT benchmark preparation, runs, and reports.""" + +import argparse +import copy +import json +import uuid +from datetime import datetime, timezone +from pathlib import Path + +from pypaimon.benchmark.act.compare import ( + compare_results, + load_result_documents, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.runner import prepare_experiment, run_experiment + + +_CONFIG_ARGUMENTS = ( + ("seed", int), + ("action_horizon", int), + ("batch_size", int), + ("optimizer_steps", int), + ("image_height", int), + ("image_width", int), + ("learning_rate", float), + ("weight_decay", float), + ("warmup_batches", int), + ("timed_batches", int), + ("fetch_batches", int), + ("rounds", int), +) + + +def main(argv=None): + """Parse an ACT benchmark subcommand and write its JSON artifact.""" + parser = _parser() + args = parser.parse_args(argv) + if args.command == "prepare": + definition = copy.deepcopy(load_experiment(args.experiment)) + for name, _ in _CONFIG_ARGUMENTS: + value = getattr(args, name) + if value is not None: + definition["config"][name] = value + for name in ( + "statistics_version", "train_episode_id", + "validation_episode_id"): + value = getattr(args, name) + if value is not None: + definition[name] = value + output = Path(args.output) + experiment = prepare_experiment( + args.input, + args.warehouse, + output, + definition=definition, + database=args.database, + ) + _print_artifact("experiment", output, experiment["schema_version"]) + return 0 + if args.command == "run": + experiment = load_experiment(args.experiment) + output = ( + Path(args.output) + if args.output else _artifact_path( + args.results_dir, + "%s-%s" % (experiment["benchmark_id"], args.backend), + ) + ) + result = run_experiment( + args.backend, + args.experiment, + output, + input_root=args.input, + warehouse=args.warehouse, + ) + _print_artifact("result", output, result["status"]) + return 0 + results_dir = args.results_dir + if not args.results and results_dir is None: + results_dir = "act-results" + results = load_result_documents(args.results, results_dir=results_dir) + comparison = compare_results(results) + output = ( + Path(args.output) + if args.output else _artifact_path( + results_dir or "act-results", "comparison") + ) + output.parent.mkdir(parents=True, exist_ok=True) + output.write_text( + json.dumps(comparison, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + _print_artifact("comparison", output, comparison["status"]) + return 0 if comparison["status"] == "SUCCEEDED" else 1 + + +def _parser(): + parser = argparse.ArgumentParser( + description=__doc__, + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + commands = parser.add_subparsers(dest="command", required=True) + + prepare = commands.add_parser( + "prepare", + help="Resolve a shared experiment against matching HDF5 and Paimon data.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + prepare.add_argument("--input", required=True, help="RoboMIND HDF5 root.") + prepare.add_argument("--warehouse", required=True, help="Paimon warehouse.") + prepare.add_argument( + "--experiment", + help="Input experiment JSON; packaged defaults are used when omitted.", + ) + prepare.add_argument( + "--output", default="act-results/experiment.json", + help="Resolved experiment JSON path.") + prepare.add_argument("--database", default="robomind") + prepare.add_argument("--statistics-version") + prepare.add_argument("--train-episode-id") + prepare.add_argument("--validation-episode-id") + for name, argument_type in _CONFIG_ARGUMENTS: + prepare.add_argument( + "--" + name.replace("_", "-"), type=argument_type, default=None) + + run = commands.add_parser( + "run", + help="Run one storage backend using a resolved experiment.", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + run.add_argument("--backend", required=True, choices=("hdf5", "paimon")) + run.add_argument("--experiment", required=True) + run.add_argument("--input", help="HDF5 root; required for backend=hdf5.") + run.add_argument( + "--warehouse", help="Paimon warehouse; required for backend=paimon.") + run.add_argument("--output", help="Explicit result JSON path.") + run.add_argument( + "--results-dir", default="act-results", + help="Directory for an automatically named result.") + + compare = commands.add_parser( + "compare", + help=( + "Group results by experiment and aggregate compatible repeats." + ), + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + compare.add_argument("results", nargs="*", help="Explicit result JSON files.") + compare.add_argument( + "--results-dir", + help="Also discover ACT result JSON files in this directory.") + compare.add_argument("--output", help="Comparison JSON path.") + return parser + + +def _artifact_path(directory, prefix): + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + return Path(directory).expanduser() / ( + "%s-%s-%s.json" % (prefix, timestamp, uuid.uuid4().hex[:8])) + + +def _print_artifact(kind, path, status): + print(json.dumps({ + "artifact": str(Path(path).expanduser().resolve()), + "kind": kind, + "status": status, + }, sort_keys=True)) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/paimon-python/pypaimon/benchmark/act/compare.py b/paimon-python/pypaimon/benchmark/act/compare.py new file mode 100644 index 000000000000..ab53f4dd6706 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/compare.py @@ -0,0 +1,231 @@ +# 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. + +"""Validate and aggregate independently produced ACT benchmark results.""" + +import hashlib +import json +from pathlib import Path + + +_METRICS = { + "batch_fetch_samples_per_s": "higher", + "dataset_build_s": "lower", + "first_batch_s": "lower", + "fixed_steps_s": "lower", + "python_peak_allocated_bytes": "lower", + "wall_time_s": "lower", +} + + +def canonical_sha256(value): + """Return the SHA-256 of a JSON value using canonical serialization.""" + payload = json.dumps(value, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def load_result_documents(paths, results_dir=None): + """Load explicit result files plus ACT results discovered in a directory. + + Explicit paths must contain result documents. Directory discovery ignores + experiment and prior comparison JSON files. A path found both ways is read + once, preserving explicit-path order followed by sorted directory entries. + """ + candidates = [Path(path).expanduser().resolve() for path in paths] + explicit = set(candidates) + if results_dir is not None: + directory = Path(results_dir).expanduser().resolve() + candidates.extend(sorted(directory.glob("*.json"))) + seen = set() + results = [] + for path in candidates: + path = path.resolve() + if path in seen: + continue + seen.add(path) + with path.open(encoding="utf-8") as result_file: + document = json.load(result_file) + if document.get("schema_version") != "act-benchmark-result@1": + if path in explicit: + raise ValueError("Not an ACT benchmark result: %s." % path) + continue + results.append(document) + if not results: + raise ValueError("No ACT benchmark result files were found.") + return results + + +def compare_results(results): + """Group result documents by experiment and compare compatible backends. + + Results from different experiment definitions remain separate. Results in + one experiment group must report the same runtime environment; otherwise + the group is marked incompatible and no performance ratios are produced. + + Args: + results: Iterable of decoded ``act-benchmark-result@1`` documents. + + Returns: + A JSON-compatible comparison document with one entry per experiment. + """ + groups = {} + for result in results: + if result.get("schema_version") != "act-benchmark-result@1": + raise ValueError("Unsupported ACT benchmark result schema.") + experiment = result.get("experiment") + if not isinstance(experiment, dict): + raise ValueError("ACT benchmark result has no experiment object.") + experiment_sha256 = canonical_sha256(experiment) + if result.get("experiment_sha256") != experiment_sha256: + raise ValueError("ACT result experiment SHA-256 differs.") + if result.get("status") != "SUCCEEDED": + raise ValueError("ACT comparison requires successful results.") + if result.get("backend") not in ("hdf5", "paimon"): + raise ValueError("ACT result has an unsupported backend.") + groups.setdefault(experiment_sha256, []).append(result) + + experiments = [ + _compare_experiment(experiment_sha256, grouped) + for experiment_sha256, grouped in sorted(groups.items()) + ] + statuses = {item["status"] for item in experiments} + if statuses == {"SUCCEEDED"}: + status = "SUCCEEDED" + elif "FAILED" in statuses: + status = "FAILED" + else: + status = "INCOMPATIBLE" + return { + "schema_version": "act-benchmark-comparison@1", + "status": status, + "experiments": experiments, + } + + +def _compare_experiment(experiment_sha256, results): + environments = { + canonical_sha256(result.get("environment", {})) for result in results + } + by_backend = {} + for result in results: + by_backend.setdefault(result["backend"], []).append(result) + if set(by_backend) != {"hdf5", "paimon"}: + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "INCOMPATIBLE", + "reason": "both hdf5 and paimon results are required", + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": {}, + } + if len(environments) != 1: + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "INCOMPATIBLE", + "reason": "runtime environments differ", + "environment_sha256s": sorted(environments), + "backends": sorted(by_backend), + "metrics": {}, + } + models = {canonical_sha256(result.get("model")) for result in results} + if len(models) != 1: + return _failed_group( + experiment_sha256, by_backend, results, "models differ") + fingerprints = { + result.get("tensor_fingerprint", {}).get("sha256") + for result in results + } + if len(fingerprints) != 1 or None in fingerprints: + return _failed_group( + experiment_sha256, + by_backend, + results, + "tensor fingerprints differ", + ) + loss_traces = {canonical_sha256([{ + "round": run["round"], + "train_loss": run["train_loss"], + "validation_loss": run["validation_loss"], + } for run in result.get("runs", [])]) for result in results} + if len(loss_traces) != 1: + return _failed_group( + experiment_sha256, by_backend, results, "loss traces differ") + + medians = { + backend: _aggregate_backend(items) + for backend, items in by_backend.items() + } + metrics = {} + for name, preferred in _METRICS.items(): + values = { + backend: summary[name] + for backend, summary in medians.items() + if name in summary + } + if values: + metric = dict(values) + metric["preferred"] = preferred + if set(values) == {"hdf5", "paimon"}: + if preferred == "higher" and values["hdf5"]: + metric["paimon_over_hdf5"] = ( + values["paimon"] / values["hdf5"]) + elif preferred == "lower" and values["paimon"]: + metric["hdf5_over_paimon"] = ( + values["hdf5"] / values["paimon"]) + metrics[name] = metric + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "SUCCEEDED", + "environment": results[0]["environment"], + "environment_sha256": next(iter(environments)), + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": metrics, + } + + +def _failed_group(experiment_sha256, by_backend, results, reason): + return { + "experiment_sha256": experiment_sha256, + "experiment": results[0]["experiment"], + "status": "FAILED", + "reason": reason, + "backends": sorted(by_backend), + "result_count": len(results), + "metrics": {}, + } + + +def _aggregate_backend(results): + names = set.intersection(*( + set(result.get("summary", {})) for result in results + )) + aggregated = {} + for name in names: + if name not in _METRICS: + continue + values = [result["summary"][name]["median"] for result in results] + values.sort() + middle = len(values) // 2 + aggregated[name] = ( + values[middle] + if len(values) % 2 + else (values[middle - 1] + values[middle]) / 2.0 + ) + return aggregated diff --git a/paimon-python/pypaimon/benchmark/act/default_experiment.json b/paimon-python/pypaimon/benchmark/act/default_experiment.json new file mode 100644 index 000000000000..df3241ca683d --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/default_experiment.json @@ -0,0 +1,22 @@ +{ + "benchmark_id": "robomind-act", + "config": { + "action_horizon": 32, + "batch_size": 2, + "fetch_batches": 8, + "image_height": 64, + "image_width": 80, + "learning_rate": 0.0001, + "optimizer_steps": 2, + "rounds": 3, + "seed": 20260825, + "timed_batches": 4, + "warmup_batches": 1, + "weight_decay": 0.0001 + }, + "dataset": "RoboMIND AgileX", + "schema_version": "act-benchmark-experiment@1", + "statistics_version": "robomind-agilex-joint-position@1", + "train_episode_id": null, + "validation_episode_id": null +} diff --git a/paimon-python/pypaimon/benchmark/act/experiment.py b/paimon-python/pypaimon/benchmark/act/experiment.py new file mode 100644 index 000000000000..f1bb675ca298 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/experiment.py @@ -0,0 +1,42 @@ +# 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. + +"""Load the declarative parameters shared by ACT benchmark runs.""" + +import json +from pathlib import Path + + +DEFAULT_EXPERIMENT = Path(__file__).with_name("default_experiment.json") + + +def load_experiment(path=None): + """Load an ACT experiment definition from JSON or the packaged default. + + Args: + path: Optional JSON path. When omitted, the packaged RoboMIND ACT + benchmark defaults are loaded. + + Returns: + A dictionary containing the benchmark identity, normalization version, + episode selection, and shared ACT/measurement configuration. + """ + source = DEFAULT_EXPERIMENT if path is None else Path(path) + with source.expanduser().open(encoding="utf-8") as experiment_file: + experiment = json.load(experiment_file) + if not isinstance(experiment, dict): + raise ValueError("ACT experiment must be a JSON object.") + return experiment diff --git a/paimon-python/pypaimon/benchmark/act_harness.py b/paimon-python/pypaimon/benchmark/act/harness.py similarity index 77% rename from paimon-python/pypaimon/benchmark/act_harness.py rename to paimon-python/pypaimon/benchmark/act/harness.py index 8b121ec5cd74..fa8164765a4f 100644 --- a/paimon-python/pypaimon/benchmark/act_harness.py +++ b/paimon-python/pypaimon/benchmark/act/harness.py @@ -42,7 +42,12 @@ @dataclass(frozen=True) class BenchmarkConfig: - """One immutable ACT and measurement configuration for both backends.""" + """Immutable model, sampling, training, and measurement parameters. + + Every backend reconstructs this configuration from the resolved experiment + so tensor shapes, optimizer behavior, random seeds, and metric boundaries + remain comparable. + """ seed: int = 20260825 action_horizon: int = 32 @@ -92,7 +97,13 @@ def to_dict(self): @dataclass(frozen=True) class WindowPlan: - """Explicit window indices consumed identically by both backends.""" + """Logical dataset-window indices consumed by one experiment. + + Measurement indices cover warm-up and timed reads, train indices cover + fixed optimizer steps, and validation indices cover the final loss. These + are map-style dataset indices, not Paimon row IDs. ``sha256`` identifies + the exact plan across independent backend processes. + """ seed: int measurement_indices: tuple @@ -115,7 +126,18 @@ def to_dict(self): def build_window_plan(train_window_count, validation_window_count, config): - """Build stable indices for training, validation, and batch-fetch timing.""" + """Build deterministic measurement, training, and validation indices. + + Args: + train_window_count: Number of complete windows in the train dataset. + validation_window_count: Number of complete validation windows. + config: Shared benchmark configuration supplying counts and the seed. + + Returns: + A :class:`WindowPlan`. When more samples are needed than a dataset + contains, consecutive seeded permutations are concatenated; sampling + does not become independent sampling with replacement. + """ train_window_count = _positive_int( train_window_count, "train_window_count") validation_window_count = _positive_int( @@ -135,15 +157,40 @@ def build_window_plan(train_window_count, validation_window_count, config): def decode_rgb_image(payload): - """Decode one JPEG/PNG payload identically for HDF5 and Paimon.""" + """Decode JPEG/PNG bytes into an ``H x W x 3`` RGB NumPy array. + + Raises: + ValueError: If Pillow cannot decode the payload as an image. + """ 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 decode_image_tensor(value): + """Decode bytes or an HDF5 uint8 value into normalized ``C x H x W``. + + The returned NumPy array is float32 with values in ``[0, 1]``. Both + storage backends call this function so image conversion is not part of the + performance difference being measured. + """ + 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 + + def validate_act_batch(batch, config): - """Validate the exact tensor contract passed to the shared ACT policy.""" + """Validate a collated batch against the shared ACT tensor contract. + + Successful validation returns ``None``. It checks exact fields, tensor + shapes and dtypes, finite values, image range, complete unpadded windows, + and ``sample_id == episode_id#step_idx`` identity. + """ required = { "sample_id", "episode_id", "step_idx", "qpos", "action", "images", "is_pad", @@ -181,7 +228,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("Paired ACT benchmark windows must be complete and unpadded.") + raise ValueError("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()): @@ -191,7 +238,11 @@ def validate_act_batch(batch, config): def build_lerobot_batch(batch, config): - """Map the common window contract to LeRobot ACTPolicy feature names.""" + """Map a shared ACT batch to LeRobot ``ACTPolicy`` feature names. + + Images are resized bilinearly to the configured height and width when + necessary. State, action, and padding retain their original semantics. + """ validate_act_batch(batch, config) images = batch["images"] target_size = (config.image_height, config.image_width) @@ -211,7 +262,14 @@ def build_lerobot_batch(batch, config): def build_act_policy(config): - """Build the reduced CPU LeRobot ACT configuration used by the benchmark.""" + """Build the reduced CPU ACT policy used only by this benchmark. + + Returns: + ``(policy, metadata)`` containing the LeRobot policy and a + JSON-compatible description of its architecture and parameter counts. + Pretrained weights are disabled, so this function performs no model + download and does not represent a production training configuration. + """ try: import importlib.metadata from lerobot.configs.types import FeatureType, PolicyFeature @@ -219,7 +277,7 @@ def build_act_policy(config): from lerobot.policies.act.modeling_act import ACTPolicy except ImportError as error: raise ImportError( - "Paired ACT benchmark requires: " + "ACT benchmark requires: " "pip install -e '.[act]'.") from error inputs = { @@ -282,7 +340,20 @@ def run_backend( config, sample_sequence_sha256, policy_factory=None): - """Measure a backend with the shared plan, model, and trainer.""" + """Measure one backend with the shared plan, model, and trainer. + + ``backend`` is a result label and ``round_number`` identifies the repeat. + ``dataset_factory`` must return ``(train_dataset, validation_dataset)`` and + must be reusable: it is called for the timed run and again by the separate + Python-memory replay. ``policy_factory`` is an optional test hook returning + ``(policy, model_metadata)``. + + Returns: + A JSON-compatible metrics dictionary covering dataset construction, + first batch, timed batch fetch, fixed optimizer steps, validation loss, + and a separate ``tracemalloc`` peak replay. OS page cache is not + controlled and native Arrow/Torch allocations are outside tracemalloc. + """ _seed_everything(config.seed) policy_factory = policy_factory or build_act_policy started = time.monotonic() @@ -408,6 +479,11 @@ def run_backend( def _measure_python_peak(dataset_factory, plan, config): + """Measure Python allocation peak in a separate dataset-first-batch replay. + + The factory is called again so tracing overhead cannot distort the main + throughput timings. The returned integer is the tracemalloc peak in bytes. + """ gc.collect() tracemalloc.start() try: @@ -428,6 +504,7 @@ def _measure_python_peak(dataset_factory, plan, config): def _repeat_permutations(size, count, seed): + """Return ``count`` indices by concatenating seeded permutations.""" values = [] generator = np.random.RandomState(seed) while len(values) < count: @@ -436,6 +513,7 @@ def _repeat_permutations(size, count, seed): def _seed_everything(seed): + """Reset Python, NumPy, and Torch RNGs and enable deterministic Torch ops.""" random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) @@ -462,6 +540,21 @@ def _positive_int(value, name): def _iter_logical_batches( dataset, indices, *, logical_batch_size, fetch_batches): + """Yield collated model batches while coalescing physical dataset reads. + + Args: + dataset: Map-style dataset implementing ``__getitem__`` and optionally + plural ``__getitems__(indices)`` access. + indices: Explicit ordered logical-window indices. Their count must be + divisible by ``logical_batch_size``. + logical_batch_size: Number of samples consumed by one model step. + fetch_batches: Logical batches combined into one physical dataset read. + + Yields: + Collated logical batches in the exact input-index order. A plural + dataset method is preferred when available; otherwise samples are read + individually and split back into the same logical batches. + """ logical_batch_size = _positive_int( logical_batch_size, "logical_batch_size") fetch_batches = _positive_int(fetch_batches, "fetch_batches") diff --git a/paimon-python/pypaimon/benchmark/act/hdf5.py b/paimon-python/pypaimon/benchmark/act/hdf5.py new file mode 100644 index 000000000000..cd46a733bad2 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/hdf5.py @@ -0,0 +1,196 @@ +# 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. + +"""HDF5 dataset adapter for the RoboMIND ACT benchmark.""" + +import numpy as np +import torch +from torch.utils.data import Dataset + +from pypaimon.benchmark.act.harness import decode_image_tensor + + +QPOS_FIELDS = ( + "puppet/joint_position_left", + "puppet/joint_position_right", +) +ACTION_FIELDS = ( + "master/joint_position_left", + "master/joint_position_right", +) +IMAGE_FIELDS = ( + "observations/rgb_images/camera_front", + "observations/rgb_images/camera_left_wrist", + "observations/rgb_images/camera_right_wrist", +) + + +class Hdf5ACTWindowDataset(Dataset): + """Read complete ACT windows lazily from one HDF5 episode. + + ``episode`` supplies the file path, logical episode ID, and frame count. + For a window anchor, state and three camera images come from the anchor + frame while actions cover ``[anchor, anchor + action_horizon)``. Each + access opens and closes the HDF5 file and returns the shared ACT sample + mapping consumed by :mod:`pypaimon.benchmark.act.harness`. + """ + + 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, QPOS_FIELDS, anchor) + action = _read_vectors( + h5, + ACTION_FIELDS, + slice(anchor, anchor + self.action_horizon), + ) + images = np.stack([ + decode_image_tensor(h5[field][anchor]) for field in 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), + } + + +def create_datasets(train_episode, validation_episode, normalization, config): + """Create HDF5 datasets for the experiment's selected episodes. + + Args: + train_episode: Selected training episode with its HDF5 path and frame + count. + validation_episode: Selected validation episode with the same fields. + normalization: Shared state and action normalization arrays. + config: Benchmark configuration containing the action horizon. + + Returns: + ``(train_dataset, validation_dataset)`` in that order. + """ + return ( + Hdf5ACTWindowDataset( + train_episode, normalization, config.action_horizon), + Hdf5ACTWindowDataset( + validation_episode, normalization, config.action_horizon), + ) + + +def compute_normalization(episodes): + """Compute train-only HDF5 state and action normalization. + + Args: + episodes: Discovered episodes carrying ``path``, ``split``, and + ``success`` attributes. + + Returns: + ``(normalization, metadata)`` where normalization contains float32 + arrays used by training. Metadata retains the float64 action moments + and frame count used to validate Paimon statistics without losing + precision. Standard deviations use a ``1e-2`` floor. + """ + 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, QPOS_FIELDS, slice(None), dtype=np.float64)) + action.update(_read_vectors( + h5, ACTION_FIELDS, slice(None), dtype=np.float64)) + qpos_mean, qpos_std = qpos.finish() + action_mean, action_std = action.finish() + return ({ + "qpos_mean": qpos_mean.astype(np.float32), + "qpos_std": qpos_std.astype(np.float32), + "action_mean": action_mean.astype(np.float32), + "action_std": action_std.astype(np.float32), + }, { + "action_mean": action_mean, + "action_std": action_std, + "frame_count": action.count, + }) + + +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 + + +class _Moments(object): + """Accumulate float64 population moments with a ``1e-2`` std floor.""" + + 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) diff --git a/paimon-python/pypaimon/benchmark/act/paimon.py b/paimon-python/pypaimon/benchmark/act/paimon.py new file mode 100644 index 000000000000..0fd3f96582a9 --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/paimon.py @@ -0,0 +1,145 @@ +# 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. + +"""Paimon dataset adapter for the RoboMIND ACT benchmark.""" + +import numpy as np +import torch + +from pypaimon.benchmark.act.harness import decode_image_tensor +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", +) + + +class PaimonACTAdapter: + """Convert a contiguous Paimon row window to the shared ACT sample. + + State and camera columns are taken from the anchor row. The action column + covers the full horizon. The returned mapping has the same IDs, tensors, + shapes, and normalization as :class:`Hdf5ACTWindowDataset`. + """ + + 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_tensor(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 create_datasets( + frames, + snapshot_id, + train_episode_id, + validation_episode_id, + normalization, + config): + """Create lazy train and validation windows pinned to one snapshot. + + Image columns are anchor-only, so one sample reads three observation + images rather than one image set per action-horizon row. + + Args: + frames: Paimon frames table used to create both scans. + snapshot_id: Snapshot pinned by experiment preparation. Both returned + datasets reject any different resolved snapshot. + train_episode_id: Episode selected for training windows. + validation_episode_id: Episode selected for validation windows. + normalization: Shared state and action normalization arrays. + config: Benchmark configuration containing the action horizon. + + Returns: + ``(train_dataset, validation_dataset)`` in that order, as lazy + ``ContiguousWindowDataset`` instances pinned to ``snapshot_id``. + """ + datasets = tuple( + frames.scan(snapshot_id=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, + anchor_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 != {snapshot_id}: + raise RuntimeError( + "Paimon ACT windows must remain pinned to frames snapshot %s; " + "got %s." % (snapshot_id, sorted(actual_snapshot_ids))) + return datasets + + +def statistics_row(connection, statistics_version): + """Return the unique versioned action-statistics row.""" + 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 latest_snapshot_id(table): + """Return the table's latest snapshot ID or fail for an empty 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 diff --git a/paimon-python/pypaimon/benchmark/act/runner.py b/paimon-python/pypaimon/benchmark/act/runner.py new file mode 100644 index 000000000000..226e0dbdac3c --- /dev/null +++ b/paimon-python/pypaimon/benchmark/act/runner.py @@ -0,0 +1,691 @@ +# 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. + +"""Prepare and run RoboMIND ACT benchmarks over HDF5 or 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. Each backend runs independently without attempting OS cache control +and writes its tensor fingerprint, loss trace, timing metrics, and Python +allocation metrics to one result JSON document. +Ingestion and canonical-action backfill are deliberately outside the benchmark. +""" + +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 +import pypaimon.multimodal as pmm +from pypaimon.benchmark.act.hdf5 import ( + compute_normalization as compute_hdf5_normalization, + create_datasets as create_hdf5_datasets, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.harness import ( + BenchmarkConfig, + WindowPlan, + build_window_plan, + run_backend, +) +from pypaimon.benchmark.act.compare import canonical_sha256 +from pypaimon.benchmark.act.paimon import ( + create_datasets as create_paimon_datasets, + latest_snapshot_id, + statistics_row, +) +from pypaimon.sample import robomind_agilex as agilex + + +@dataclass(frozen=True) +class _BenchmarkEpisode: + path: Path + source_key: str + episode_id: str + split: str + success: bool + frame_count: int + + +def prepare_experiment( + input_root, + warehouse, + output_path, + *, + definition=None, + database=agilex.DEFAULT_DATABASE): + """Resolve a benchmark definition against matching HDF5 and Paimon data. + + Preparation is outside timed benchmark execution. It verifies source + identity and Paimon statistics, selects eligible train/validation episodes, + computes train-only normalization, and fixes every logical window index. + + Args: + input_root: RoboMIND AgileX HDF5 root used as the source of episode + files and raw normalization moments. + warehouse: Existing Paimon warehouse containing the matching ingested + and canonical-action-backfilled dataset. + output_path: Destination for the resolved experiment JSON document. + definition: Optional decoded experiment definition. The packaged + defaults are used when omitted. + database: Paimon database containing the RoboMIND tables. + + Returns: + The resolved, JSON-compatible experiment dictionary written to + ``output_path``. + """ + definition = load_experiment() if definition is None else definition + if definition.get("schema_version") != "act-benchmark-experiment@1": + raise ValueError("Unsupported ACT benchmark experiment schema.") + config = BenchmarkConfig(**definition["config"]) + statistics_version = definition["statistics_version"] + input_root = Path(input_root).expanduser().resolve() + warehouse = Path(warehouse).expanduser().resolve() + output_path = Path(output_path).expanduser().resolve() + + discovered = agilex.discover_episodes(input_root) + connection = pmm.connect( + database=database, options={"warehouse": str(warehouse)}) + source_episodes, source_sha256 = _validate_source_identity( + discovered, _episode_rows(connection)) + source_by_id = {episode.episode_id: episode for episode in source_episodes} + frames = connection.get_table(agilex.FRAMES_TABLE) + frames_snapshot_id = latest_snapshot_id(frames) + normalization, normalization_metadata = _shared_normalization( + source_episodes, + connection, + frames_snapshot_id, + statistics_version, + ) + del normalization + train_episode = _select_episode( + source_by_id, + split="train", + requested=definition.get("train_episode_id"), + action_horizon=config.action_horizon, + ) + validation_episode = _select_episode( + source_by_id, + split="val", + requested=definition.get("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) + episodes = sorted(({ + "episode_id": episode.episode_id, + "source_key": episode.source_key, + "split": episode.split, + "success": episode.success, + "frame_count": episode.frame_count, + } for episode in source_episodes), key=lambda item: item["episode_id"]) + experiment = { + "schema_version": "act-benchmark-experiment@1", + "benchmark_id": definition.get("benchmark_id", "robomind-act"), + "dataset": definition.get("dataset", "RoboMIND AgileX"), + "config": config.to_dict(), + "statistics_version": statistics_version, + "train_episode_id": train_episode.episode_id, + "validation_episode_id": validation_episode.episode_id, + "source": { + "sha256": source_sha256, + "episodes": episodes, + }, + "normalization": normalization_metadata, + "window_plan": { + **plan.to_dict(), + "sha256": plan.sha256, + "sample_sequence_sha256": sequence_sha256, + }, + "paimon": { + "database": database, + "frames_table": agilex.FRAMES_TABLE, + "frames_snapshot_id": frames_snapshot_id, + }, + } + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(experiment, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return experiment + + +def run_experiment( + backend, + experiment_path, + output_path, + *, + input_root=None, + warehouse=None, + policy_factory=None): + """Run one storage backend against a resolved ACT experiment. + + Args: + backend: Result label and dataset implementation, either ``hdf5`` or + ``paimon``. + experiment_path: Resolved JSON produced by :func:`prepare_experiment`. + output_path: Destination JSON result path. + input_root: Required only for the HDF5 backend. + warehouse: Required only for the Paimon backend. + policy_factory: Optional test hook returning ``(policy, metadata)``. + + Returns: + A JSON-compatible single-backend result containing the resolved + experiment, runtime environment, tensor fingerprint, per-round raw + metrics, and median/min/max summary. + """ + if backend not in ("hdf5", "paimon"): + raise ValueError("backend must be 'hdf5' or 'paimon'.") + experiment = load_experiment(experiment_path) + _validate_resolved_experiment(experiment) + config = BenchmarkConfig(**experiment["config"]) + plan = _window_plan_from_experiment(experiment) + normalization = { + name: np.asarray(value, dtype=np.float32) + for name, value in experiment["normalization"]["values"].items() + } + sequence_sha256 = experiment["window_plan"]["sample_sequence_sha256"] + if backend == "hdf5": + if input_root is None: + raise ValueError("input_root is required for the HDF5 backend.") + episodes = _hdf5_episodes_from_experiment(input_root, experiment) + by_id = {episode.episode_id: episode for episode in episodes} + train_episode = by_id[experiment["train_episode_id"]] + validation_episode = by_id[experiment["validation_episode_id"]] + + def dataset_factory(): + return create_hdf5_datasets( + train_episode, validation_episode, normalization, config) + + source = {"input_root": str(Path(input_root).expanduser().resolve())} + else: + if warehouse is None: + raise ValueError("warehouse is required for the Paimon backend.") + dataset_factory, source = _paimon_factory_from_experiment( + warehouse, experiment, normalization, config) + + started_at = _utc_now() + started = time.monotonic() + fingerprint = _tensor_fingerprint(dataset_factory(), plan) + runs = [] + for round_number in range(1, config.rounds + 1): + runs.append(run_backend( + backend, + round_number, + dataset_factory, + plan, + config, + sequence_sha256, + policy_factory=policy_factory, + )) + gc.collect() + result = { + "schema_version": "act-benchmark-result@1", + "benchmark_id": experiment["benchmark_id"], + "run_id": "%s-%s" % ( + started_at.replace(":", "").replace("-", ""), + uuid.uuid4().hex[:8], + ), + "status": "SUCCEEDED", + "backend": backend, + "experiment": experiment, + "experiment_sha256": canonical_sha256(experiment), + "source": source, + "tensor_fingerprint": fingerprint, + "model": runs[0]["model"], + "runs": runs, + "summary": _summarize(runs), + "environment": { + "python": platform.python_version(), + "os": platform.platform(), + "machine": platform.machine(), + "torch": torch.__version__, + "source_commit": _git_head(Path(__file__).resolve().parents[4]), + }, + "command": _command_argv(), + "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 dataset loading, distributed training, and " + "recovery are unverified.", + "Python tracemalloc excludes native Arrow and Torch allocations.", + ], + "started_at": started_at, + "finished_at": _utc_now(), + } + output_path = Path(output_path).expanduser().resolve() + output_path.parent.mkdir(parents=True, exist_ok=True) + output_path.write_text( + json.dumps(result, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return result + + +def _validate_resolved_experiment(experiment): + """Reject incomplete or internally inconsistent resolved experiments.""" + required = { + "schema_version", "benchmark_id", "dataset", "config", + "statistics_version", "train_episode_id", "validation_episode_id", + "source", "normalization", "window_plan", "paimon", + } + if experiment.get("schema_version") != "act-benchmark-experiment@1": + raise ValueError("Unsupported ACT benchmark experiment schema.") + missing = required - set(experiment) + if missing: + raise ValueError( + "Resolved ACT experiment is missing: %s." + % ", ".join(sorted(missing))) + source = experiment["source"] + if canonical_sha256(source["episodes"]) != source["sha256"]: + raise ValueError("ACT experiment source-manifest hash differs.") + normalization = experiment["normalization"] + if canonical_sha256(normalization["values"]) != normalization["sha256"]: + raise ValueError("ACT experiment normalization hash differs.") + plan = _window_plan_from_experiment(experiment) + if plan.sha256 != experiment["window_plan"]["sha256"]: + raise ValueError("ACT experiment window-plan hash differs.") + episodes = { + item["episode_id"]: item for item in source["episodes"] + } + try: + train = episodes[experiment["train_episode_id"]] + validation = episodes[experiment["validation_episode_id"]] + except KeyError as error: + raise ValueError( + "ACT experiment selected episode is absent from the source." + ) from error + config = BenchmarkConfig(**experiment["config"]) + expected_plan = build_window_plan( + train["frame_count"] - config.action_horizon + 1, + validation["frame_count"] - config.action_horizon + 1, + config, + ) + if expected_plan.to_dict() != plan.to_dict(): + raise ValueError( + "ACT experiment window plan was not built from its config and " + "selected episodes.") + expected_sequence = _sample_sequence_sha256( + train["episode_id"], validation["episode_id"], plan) + if expected_sequence != experiment["window_plan"][ + "sample_sequence_sha256"]: + raise ValueError("ACT experiment sample-sequence hash differs.") + + +def _window_plan_from_experiment(experiment): + """Reconstruct immutable logical-window indices from JSON values.""" + value = experiment["window_plan"] + return WindowPlan( + seed=value["seed"], + measurement_indices=tuple(value["measurement_indices"]), + train_indices=tuple(value["train_indices"]), + validation_indices=tuple(value["validation_indices"]), + ) + + +def _hdf5_episodes_from_experiment(input_root, experiment): + """Validate HDF5 episode identity and attach manifest frame counts.""" + discovered = agilex.discover_episodes(Path(input_root).expanduser().resolve()) + by_id = {episode.episode_id: episode for episode in discovered} + expected = experiment["source"]["episodes"] + actual_identity = sorted(({ + "episode_id": episode.episode_id, + "source_key": episode.source_key, + "split": episode.split, + "success": episode.success, + } for episode in discovered), key=lambda item: item["episode_id"]) + expected_identity = [{ + "episode_id": item["episode_id"], + "source_key": item["source_key"], + "split": item["split"], + "success": item["success"], + } for item in expected] + if actual_identity != expected_identity: + raise ValueError("HDF5 source differs from the ACT experiment.") + return [ + _BenchmarkEpisode( + path=by_id[item["episode_id"]].path, + source_key=item["source_key"], + episode_id=item["episode_id"], + split=item["split"], + success=item["success"], + frame_count=item["frame_count"], + ) + for item in expected + ] + + +def _paimon_factory_from_experiment( + warehouse, experiment, normalization, config): + """Validate Paimon source/statistics and return a pinned dataset factory.""" + warehouse = Path(warehouse).expanduser().resolve() + paimon = experiment["paimon"] + connection = pmm.connect( + database=paimon["database"], + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(paimon["frames_table"]) + snapshot_id = paimon["frames_snapshot_id"] + expected_episodes = experiment["source"]["episodes"] + actual_episodes = sorted(_episode_rows(connection), + key=lambda item: item["episode_id"]) + if actual_episodes != expected_episodes: + raise ValueError("Paimon source differs from the ACT experiment.") + row = statistics_row(connection, experiment["statistics_version"]) + expected_normalization = experiment["normalization"] + action_mean = np.asarray(row["action_mean"], dtype=np.float32) + action_std = np.asarray(row["action_std"], dtype=np.float32) + if ( + row["source_snapshot_id"] != snapshot_id + or row["source_split"] != "train" + or row["frame_count"] != expected_normalization["frame_count"] + or row["feature_name"] != "action" + or row["standard_deviation_floor"] != 1e-2 + or not np.array_equal( + action_mean, normalization["action_mean"]) + or not np.array_equal(action_std, normalization["action_std"])): + raise ValueError( + "Paimon normalization differs from the ACT experiment.") + def factory(): + return create_paimon_datasets( + frames, + snapshot_id, + experiment["train_episode_id"], + experiment["validation_episode_id"], + normalization, + config, + ) + + return factory, { + "warehouse": str(warehouse), + "database": paimon["database"], + "frames_table": paimon["frames_table"], + "frames_snapshot_id": snapshot_id, + } + + +def _tensor_fingerprint(datasets, plan): + """Hash the exact planned sample IDs and tensors outside timed execution.""" + comparisons = ( + ("train", datasets[0], + sorted(set(plan.measurement_indices + plan.train_indices))), + ("validation", datasets[1], + sorted(set(plan.validation_indices))), + ) + digest = hashlib.sha256() + count = 0 + for split, dataset, indices in comparisons: + for index in indices: + sample = dataset[index] + identity = { + "split": split, + "index": index, + "sample_id": sample["sample_id"], + "episode_id": sample["episode_id"], + "step_idx": sample["step_idx"], + } + digest.update(json.dumps( + identity, sort_keys=True, separators=(",", ":") + ).encode("utf-8")) + for name in ("qpos", "action", "images", "is_pad"): + tensor = sample[name].detach().cpu().contiguous() + digest.update(name.encode("utf-8")) + digest.update(str(tensor.dtype).encode("ascii")) + digest.update(str(tuple(tensor.shape)).encode("ascii")) + digest.update(tensor.numpy().tobytes()) + count += 1 + return { + "sha256": digest.hexdigest(), + "checked_window_count": count, + "fields": [ + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + ], + } + + + +def _shared_normalization( + episodes, + connection, + frames_snapshot_id, + statistics_version): + """Build one train-only normalization contract for both backends. + + HDF5 supplies state and action moments from successful train episodes. + Versioned Paimon action statistics must match the float64 HDF5 moments, + train scope, frame count, source snapshot, feature name, and ``1e-2`` + standard-deviation floor. + + Returns: + ``(arrays, metadata)`` where arrays are float32 training values and + metadata is JSON-compatible and includes their canonical SHA-256. + """ + normalization, hdf5_metadata = compute_hdf5_normalization(episodes) + action_mean = hdf5_metadata["action_mean"] + action_std = hdf5_metadata["action_std"] + action_count = hdf5_metadata["frame_count"] + 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["action_mean"] = stored_mean.astype(np.float32) + normalization["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 _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): + """Match HDF5 discovery to Paimon episodes and return a manifest hash. + + Episode ID, source key, split, and success must match exactly. Paimon's + versioned episode rows contribute frame counts used to build complete + windows. The returned records retain the local HDF5 paths while the hash + covers only portable source metadata. + """ + 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): + """Select a successful, split-matching episode long enough for one window. + + An explicit episode is honored when eligible; otherwise the + lexicographically first eligible episode ID is selected. + """ + 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 _summarize(runs): + """Return median, minimum, and maximum metrics across backend repeats.""" + metrics = ( + "dataset_build_s", + "first_batch_s", + "batch_fetch_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): + """Hash episode-qualified sample IDs in measurement/train/validation order.""" + value = { + "batch_fetch": [ + "%s#%d" % (train_episode_id, index) + for index in plan.measurement_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 _git_head(repository): + 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 _command_argv(): + """Return the invoked Python basename and command-line arguments.""" + 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") diff --git a/paimon-python/pypaimon/benchmark/paired_act.py b/paimon-python/pypaimon/benchmark/paired_act.py deleted file mode 100644 index b2bb30ebddd6..000000000000 --- a/paimon-python/pypaimon/benchmark/paired_act.py +++ /dev/null @@ -1,865 +0,0 @@ -# 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_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", - }, - "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 dataset loading, 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, - anchor_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.measurement_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", - "batch_fetch_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 = { - "batch_fetch": [ - "%s#%d" % (train_episode_id, index) - for index in plan.measurement_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_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): - 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(): - 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__, - 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( - "--action-horizon", type=int, default=BenchmarkConfig.action_horizon, - help="Number of contiguous action rows in each sample.") - parser.add_argument( - "--batch-size", type=int, default=BenchmarkConfig.batch_size, - help="Shared logical batch size.") - parser.add_argument( - "--optimizer-steps", type=int, default=BenchmarkConfig.optimizer_steps, - help="Fixed optimizer steps per backend run.") - parser.add_argument( - "--image-height", type=int, default=BenchmarkConfig.image_height, - help="ACT input image height after resizing.") - parser.add_argument( - "--image-width", type=int, default=BenchmarkConfig.image_width, - help="ACT input image width after resizing.") - parser.add_argument( - "--learning-rate", type=float, default=BenchmarkConfig.learning_rate, - help="Shared AdamW learning rate.") - parser.add_argument( - "--weight-decay", type=float, default=BenchmarkConfig.weight_decay, - help="Shared AdamW weight decay.") - parser.add_argument( - "--warmup-batches", type=int, default=BenchmarkConfig.warmup_batches, - help="Logical batches consumed before batch-fetch timing.") - parser.add_argument( - "--timed-batches", type=int, - default=BenchmarkConfig.timed_batches, - help="Logical batches used for dataset batch-fetch throughput.") - parser.add_argument( - "--fetch-batches", type=int, default=BenchmarkConfig.fetch_batches, - help="Logical batches coalesced into one physical dataset fetch.") - parser.add_argument( - "--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, - 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, - timed_batches=args.timed_batches, - fetch_batches=args.fetch_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/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py new file mode 100644 index 000000000000..81d55bc1708c --- /dev/null +++ b/paimon-python/pypaimon/tests/act_benchmark_test.py @@ -0,0 +1,233 @@ +# 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 pytest + +from pypaimon.benchmark.act.compare import ( + canonical_sha256, + compare_results, + load_result_documents, +) +from pypaimon.benchmark.act.experiment import load_experiment + + +def test_default_experiment_loads_packaged_benchmark_parameters(): + experiment = load_experiment() + + assert experiment["schema_version"] == "act-benchmark-experiment@1" + assert experiment["benchmark_id"] == "robomind-act" + assert experiment["config"]["seed"] == 20260825 + assert experiment["config"]["action_horizon"] == 32 + assert experiment["config"]["rounds"] == 3 + assert experiment["statistics_version"] == ( + "robomind-agilex-joint-position@1") + + +def test_compare_reports_ratio_for_matching_backend_results(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + results = [ + _result("hdf5", experiment, environment, throughput=10.0), + _result("paimon", experiment, environment, throughput=15.0), + ] + + comparison = compare_results(results) + + assert comparison["status"] == "SUCCEEDED" + assert len(comparison["experiments"]) == 1 + group = comparison["experiments"][0] + assert group["backends"] == ["hdf5", "paimon"] + assert group["metrics"]["batch_fetch_samples_per_s"] == { + "hdf5": 10.0, + "paimon": 15.0, + "paimon_over_hdf5": 1.5, + "preferred": "higher", + } + + +def test_compare_rejects_different_tensor_fingerprints(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5 = _result("hdf5", experiment, environment, throughput=10.0) + paimon = _result("paimon", experiment, environment, throughput=15.0) + paimon["tensor_fingerprint"]["sha256"] = "different" + + comparison = compare_results([hdf5, paimon]) + + assert comparison["status"] == "FAILED" + group = comparison["experiments"][0] + assert group["status"] == "FAILED" + assert group["reason"] == "tensor fingerprints differ" + assert group["metrics"] == {} + + +def test_compare_requires_results_from_both_backends(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + + comparison = compare_results([ + _result("hdf5", experiment, environment, throughput=10.0), + ]) + + assert comparison["status"] == "INCOMPATIBLE" + group = comparison["experiments"][0] + assert group["status"] == "INCOMPATIBLE" + assert group["reason"] == "both hdf5 and paimon results are required" + assert group["metrics"] == {} + + +def test_compare_reports_lower_is_better_metric_as_paimon_speedup(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5 = _result("hdf5", experiment, environment, throughput=10.0) + paimon = _result("paimon", experiment, environment, throughput=15.0) + hdf5["summary"]["first_batch_s"] = { + "median": 2.0, "min": 2.0, "max": 2.0} + paimon["summary"]["first_batch_s"] = { + "median": 1.0, "min": 1.0, "max": 1.0} + + comparison = compare_results([hdf5, paimon]) + + assert comparison["experiments"][0]["metrics"]["first_batch_s"] == { + "hdf5": 2.0, + "paimon": 1.0, + "hdf5_over_paimon": 2.0, + "preferred": "lower", + } + + +def test_load_results_combines_explicit_files_and_directory(tmp_path): + experiment = {"schema_version": "act-benchmark-experiment@1"} + environment = {"python": "3.10", "machine": "arm64"} + hdf5_path = tmp_path / "hdf5.json" + paimon_path = tmp_path / "paimon.json" + ignored_path = tmp_path / "experiment.json" + hdf5_path.write_text(json.dumps( + _result("hdf5", experiment, environment, throughput=10.0))) + paimon_path.write_text(json.dumps( + _result("paimon", experiment, environment, throughput=15.0))) + ignored_path.write_text(json.dumps(experiment)) + + results = load_result_documents( + [hdf5_path], results_dir=tmp_path) + + assert [result["backend"] for result in results] == ["hdf5", "paimon"] + + +def test_compare_groups_multiple_experiments_without_cross_comparing(): + environment = {"python": "3.10", "machine": "arm64"} + results = [] + for seed in (1, 2): + experiment = { + "schema_version": "act-benchmark-experiment@1", + "config": {"seed": seed}, + } + results.extend([ + _result("hdf5", experiment, environment, throughput=10.0), + _result("paimon", experiment, environment, throughput=15.0), + ]) + + comparison = compare_results(results) + + assert comparison["status"] == "SUCCEEDED" + assert len(comparison["experiments"]) == 2 + assert all( + group["backends"] == ["hdf5", "paimon"] + for group in comparison["experiments"] + ) + + +def test_compare_reports_incompatible_runtime_environments(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + hdf5 = _result( + "hdf5", experiment, + {"python": "3.10", "machine": "arm64"}, throughput=10.0) + paimon = _result( + "paimon", experiment, + {"python": "3.11", "machine": "arm64"}, throughput=15.0) + + comparison = compare_results([hdf5, paimon]) + + assert comparison["status"] == "INCOMPATIBLE" + group = comparison["experiments"][0] + assert group["reason"] == "runtime environments differ" + assert len(group["environment_sha256s"]) == 2 + assert group["metrics"] == {} + + +def test_compare_rejects_tampered_result_experiment_hash(): + experiment = {"schema_version": "act-benchmark-experiment@1"} + result = _result( + "hdf5", + experiment, + {"python": "3.10", "machine": "arm64"}, + throughput=10.0, + ) + result["experiment_sha256"] = "tampered" + + with pytest.raises(ValueError, match="experiment SHA-256 differs"): + compare_results([result]) + + +def test_shared_harness_is_imported_from_act_package(): + from pypaimon.benchmark.act.harness import BenchmarkConfig + + assert BenchmarkConfig().to_dict() == load_experiment()["config"] + + +def test_hdf5_dataset_is_owned_by_hdf5_backend(): + from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset + + assert Hdf5ACTWindowDataset.__module__ == ( + "pypaimon.benchmark.act.hdf5") + + +def test_paimon_adapter_is_owned_by_paimon_backend(): + from pypaimon.benchmark.act.paimon import PaimonACTAdapter + + assert PaimonACTAdapter.__module__ == ( + "pypaimon.benchmark.act.paimon") + + +def _result(backend, experiment, environment, throughput): + return { + "schema_version": "act-benchmark-result@1", + "status": "SUCCEEDED", + "backend": backend, + "experiment": experiment, + "experiment_sha256": canonical_sha256(experiment), + "environment": environment, + "model": {"implementation": "test-policy", "parameter_count": 1}, + "tensor_fingerprint": { + "sha256": "same-tensors", + "checked_window_count": 2, + }, + "runs": [{ + "round": round_number, + "train_loss": [1.0, 0.5], + "validation_loss": 0.25, + } for round_number in range(1, 4)], + "summary": { + "round_count": 3, + "batch_fetch_samples_per_s": { + "median": throughput, + "min": throughput, + "max": throughput, + }, + }, + } diff --git a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py b/paimon-python/pypaimon/tests/act_runner_test.py similarity index 50% rename from paimon-python/pypaimon/tests/paired_act_benchmark_test.py rename to paimon-python/pypaimon/tests/act_runner_test.py index 6e4e96002003..e7f9c0a69349 100644 --- a/paimon-python/pypaimon/tests/paired_act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -30,18 +30,22 @@ h5py = pytest.importorskip("h5py") import pypaimon.multimodal as pmm -import pypaimon.benchmark.act_harness as act_harness -import pypaimon.benchmark.paired_act as paired_act -from pypaimon.benchmark.paired_act import ( - IMAGE_COLUMNS, +import pypaimon.benchmark.act.harness as act_harness +import pypaimon.benchmark.act.__main__ as act_cli +from pypaimon.benchmark.act.runner import ( BenchmarkConfig, _git_head, - _paimon_datasets, - _shared_normalization, - _snapshot_id, - run, + prepare_experiment, + run_experiment, +) +from pypaimon.benchmark.act.experiment import load_experiment +from pypaimon.benchmark.act.compare import canonical_sha256, compare_results +from pypaimon.benchmark.act.harness import build_window_plan, run_backend +from pypaimon.benchmark.act.paimon import ( + IMAGE_COLUMNS, + create_datasets as create_paimon_datasets, + latest_snapshot_id, ) -from pypaimon.benchmark.act_harness import 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 @@ -225,7 +229,7 @@ def _write_episode(root, split, name, offset, frames=6): @pytest.fixture -def paired_input(tmp_path, monkeypatch): +def benchmark_input(tmp_path, monkeypatch): root = tmp_path / "input" _write_episode(root, "train", "train-a", 1) _write_episode(root, "train", "train-b", 11) @@ -237,7 +241,7 @@ def paired_input(tmp_path, monkeypatch): }) agilex.ingest_local(root, warehouse, batch_size=2) agilex.backfill_canonical_action( - warehouse, statistics_version="paired-test@1") + warehouse, statistics_version="act-test@1") return root, warehouse @@ -265,93 +269,220 @@ def _policy_factory(config): } -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, - timed_batches=2, - rounds=3, +def test_prepare_writes_resolved_experiment(benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "seed": 17, + "action_horizon": 3, + "batch_size": 2, + "optimizer_steps": 2, + "image_height": 8, + "image_width": 10, + "timed_batches": 2, + }) + output = tmp_path / "experiment.json" + + experiment = prepare_experiment( + input_root, warehouse, output, definition=definition) + + assert json.loads(output.read_text()) == experiment + assert experiment["schema_version"] == "act-benchmark-experiment@1" + assert experiment["train_episode_id"] == "train-a" + assert experiment["validation_episode_id"] == "val-a" + assert experiment["source"]["episodes"][0]["frame_count"] == 6 + assert len(experiment["source"]["sha256"]) == 64 + assert len(experiment["normalization"]["sha256"]) == 64 + assert len(experiment["window_plan"]["sha256"]) == 64 + assert experiment["paimon"]["frames_snapshot_id"] > 0 + + +def test_hdf5_run_consumes_resolved_experiment_without_warehouse( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "seed": 17, + "action_horizon": 3, + "batch_size": 2, + "optimizer_steps": 2, + "image_height": 8, + "image_width": 10, + "timed_batches": 2, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + result_path = tmp_path / "hdf5-result.json" + + result = run_experiment( + "hdf5", + experiment_path, + result_path, + input_root=input_root, + policy_factory=_policy_factory, ) - report = run( - input_root, - warehouse, - report_path, - config=config, - statistics_version="paired-test@1", + assert json.loads(result_path.read_text()) == result + assert result["schema_version"] == "act-benchmark-result@1" + assert result["status"] == "SUCCEEDED" + assert result["backend"] == "hdf5" + assert result["experiment"] == experiment + assert len(result["experiment_sha256"]) == 64 + assert len(result["tensor_fingerprint"]["sha256"]) == 64 + assert len(result["runs"]) == 3 + assert result["summary"]["round_count"] == 3 + + +def test_independent_backend_results_preserve_tensor_and_loss_parity( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "seed": 17, + "action_horizon": 3, + "batch_size": 2, + "optimizer_steps": 2, + "image_height": 8, + "image_width": 10, + "timed_batches": 2, + }) + experiment_path = tmp_path / "experiment.json" + prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + + hdf5_result = run_experiment( + "hdf5", + experiment_path, + tmp_path / "hdf5-result.json", + input_root=input_root, + policy_factory=_policy_factory, + ) + paimon_result = run_experiment( + "paimon", + experiment_path, + tmp_path / "paimon-result.json", + warehouse=warehouse, 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 hdf5_result["tensor_fingerprint"] == ( + paimon_result["tensor_fingerprint"]) + assert [run["train_loss"] for run in hdf5_result["runs"]] == [ + run["train_loss"] for run in paimon_result["runs"] ] - 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", - "batch_fetch_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"]) + assert [run["validation_loss"] for run in hdf5_result["runs"]] == [ + run["validation_loss"] for run in paimon_result["runs"] + ] + comparison = compare_results([hdf5_result, paimon_result]) + assert comparison["status"] == "SUCCEEDED" + assert comparison["experiments"][0]["backends"] == ["hdf5", "paimon"] + + +def test_paimon_run_rejects_normalization_not_recorded_in_statistics( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["normalization"]["values"]["action_mean"][0] += 1 + experiment["normalization"]["sha256"] = canonical_sha256( + experiment["normalization"]["values"]) + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="normalization differs"): + run_experiment( + "paimon", + experiment_path, + tmp_path / "must-not-exist.json", + warehouse=warehouse, + policy_factory=_policy_factory, + ) + + +def test_run_rejects_tampered_source_manifest(benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["source"]["episodes"][0]["frame_count"] += 1 + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="source-manifest hash differs"): + run_experiment( + "hdf5", + experiment_path, + tmp_path / "must-not-exist.json", + input_root=input_root, + policy_factory=_policy_factory, + ) + + +def test_run_rejects_config_not_used_to_build_window_plan( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + experiment = prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + experiment["config"]["seed"] += 1 + experiment_path.write_text(json.dumps(experiment)) + + with pytest.raises(ValueError, match="window plan was not built"): + run_experiment( + "hdf5", + experiment_path, + tmp_path / "must-not-exist.json", + input_root=input_root, + policy_factory=_policy_factory, + ) def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( - paired_input): - input_root, warehouse = paired_input + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"]["action_horizon"] = 3 + experiment = prepare_experiment( + input_root, + warehouse, + tmp_path / "experiment.json", + definition=definition, + ) + normalization = { + name: np.asarray(value, dtype=np.float32) + for name, value in experiment["normalization"]["values"].items() + } connection = pmm.connect( database=agilex.DEFAULT_DATABASE, options={"warehouse": str(warehouse)}, @@ -359,18 +490,12 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( frames = connection.get_table(agilex.FRAMES_TABLE) assert frames.raw_table.table_schema.options["vector.file.format"] == ( "parquet") - snapshot_id = _snapshot_id(frames) - normalization, _ = _shared_normalization( - agilex.discover_episodes(input_root), - connection, - snapshot_id, - "paired-test@1", - ) + snapshot_id = latest_snapshot_id(frames) original = ScanQuery._fetch_bodies with patch.object( ScanQuery, "_fetch_bodies", side_effect=original) as fetch: - train, validation = _paimon_datasets( + train, validation = create_paimon_datasets( frames, snapshot_id, "train-a", @@ -406,7 +531,7 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( assert train.snapshot_id == snapshot_id assert validation.snapshot_id == snapshot_id - assert _snapshot_id(frames) != snapshot_id + assert latest_snapshot_id(frames) != snapshot_id assert len(train) == 4 sample_after_append = train[0] for name in ("qpos", "action", "images", "is_pad"): @@ -414,33 +539,48 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( 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 +def test_compare_rejects_different_hdf5_tensor_bytes( + benchmark_input, tmp_path): + input_root, warehouse = benchmark_input + definition = load_experiment() + definition["statistics_version"] = "act-test@1" + definition["config"].update({ + "action_horizon": 3, + "batch_size": 1, + "optimizer_steps": 1, + "image_height": 8, + "image_width": 10, + }) + experiment_path = tmp_path / "experiment.json" + prepare_experiment( + input_root, warehouse, experiment_path, definition=definition) + paimon_result = run_experiment( + "paimon", + experiment_path, + tmp_path / "paimon-result.json", + warehouse=warehouse, + policy_factory=_policy_factory, + ) 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, - ) + hdf5_result = run_experiment( + "hdf5", + experiment_path, + tmp_path / "hdf5-result.json", + input_root=input_root, + policy_factory=_policy_factory, + ) + + comparison = compare_results([hdf5_result, paimon_result]) + assert comparison["status"] == "FAILED" + assert comparison["experiments"][0]["reason"] == ( + "tensor fingerprints differ") -def test_requires_at_least_three_alternating_rounds(): +def test_requires_at_least_three_measurement_rounds(): with pytest.raises(ValueError, match="rounds must be at least 3"): BenchmarkConfig(rounds=2) @@ -451,15 +591,42 @@ def test_fetch_batches_must_be_positive(): BenchmarkConfig(fetch_batches=0) -def test_cli_documents_physical_fetch_batches(capsys): +def test_cli_exposes_prepare_run_and_compare_contracts(capsys): + with pytest.raises(SystemExit): + act_cli.main(["prepare", "--help"]) + + prepare_help = capsys.readouterr().out + assert "--experiment" in prepare_help + assert "--fetch-batches" in prepare_help + + with pytest.raises(SystemExit): + act_cli.main(["run", "--help"]) + + run_help = capsys.readouterr().out + assert "--backend" in run_help + assert "--experiment" in run_help + assert "--results-dir" in run_help + with pytest.raises(SystemExit): - paired_act.main(["--help"]) + act_cli.main(["compare", "--help"]) + + assert "--results-dir" in capsys.readouterr().out + + +def test_automatic_artifact_paths_do_not_overwrite_same_second(tmp_path): + with patch.object(act_cli, "datetime") as now: + now.now.return_value.strftime.return_value = "20260901T120000Z" + + first = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") + second = act_cli._artifact_path(tmp_path, "robomind-act-hdf5") - assert "--fetch-batches" in capsys.readouterr().out + assert first != second + assert first.parent == tmp_path + assert second.parent == tmp_path def test_source_commit_falls_back_outside_git_checkout(tmp_path): with patch( - "pypaimon.benchmark.paired_act.subprocess.check_output", + "pypaimon.benchmark.act.runner.subprocess.check_output", side_effect=FileNotFoundError): assert _git_head(tmp_path) == "UNKNOWN" diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 382371733ab7..e2cfcc33a8af 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -233,7 +233,12 @@ def read_requirements(): version=VERSION, packages=PACKAGES, include_package_data=True, - package_data={"pypaimon": ["_full_version"]}, + package_data={ + "pypaimon": [ + "_full_version", + "benchmark/act/default_experiment.json", + ], + }, cmdclass={"build_py": PaimonBuildPy, "sdist": PaimonSdist}, install_requires=install_requires, entry_points={ From 058129e65764c000a4c90e77e7d1b710d85ec4ff Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 00:44:01 +0800 Subject: [PATCH 09/16] perf(python): read ACT state from window anchor Avoid reading the full qpos horizon when ACT only consumes the anchor observation. Add a golden backend sample contract and explicit HDF5 index bounds coverage. Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex AI-Contributed/Feature: 13/13 AI-Contributed/UT: 99/99 --- .../pypaimon/benchmark/act/paimon.py | 13 ++- .../pypaimon/tests/act_runner_test.py | 99 ++++++++++++++++++- 2 files changed, 108 insertions(+), 4 deletions(-) diff --git a/paimon-python/pypaimon/benchmark/act/paimon.py b/paimon-python/pypaimon/benchmark/act/paimon.py index 0fd3f96582a9..d99ea4227aae 100644 --- a/paimon-python/pypaimon/benchmark/act/paimon.py +++ b/paimon-python/pypaimon/benchmark/act/paimon.py @@ -47,6 +47,12 @@ def __init__(self, normalization): self.normalization = normalization def __call__(self, sample): + """Convert the generic window mapping into ACT tensors and identity. + + The persisted ``frame_index`` becomes the shared ACT ``step_idx``. + State and image columns are singleton lists; action retains the full + horizon and ``is_pad`` is forwarded unchanged. + """ qpos = np.concatenate([ np.asarray(sample[name][0], dtype=np.float32) for name in QPOS_COLUMNS @@ -86,8 +92,9 @@ def create_datasets( config): """Create lazy train and validation windows pinned to one snapshot. - Image columns are anchor-only, so one sample reads three observation - images rather than one image set per action-horizon row. + State and image columns are anchor-only, so one sample reads the initial + joint position and three observation images once rather than once per + action-horizon row. Args: frames: Paimon frames table used to create both scans. @@ -108,7 +115,7 @@ def create_datasets( ).to_contiguous_window_dataset( window_size=config.action_horizon, columns=QPOS_COLUMNS + ACTION_COLUMNS + IMAGE_COLUMNS, - anchor_columns=IMAGE_COLUMNS, + anchor_columns=QPOS_COLUMNS + IMAGE_COLUMNS, group_key="episode_id", order_key="frame_index", stride=1, diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py index e7f9c0a69349..67e1d88fe9ed 100644 --- a/paimon-python/pypaimon/tests/act_runner_test.py +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -19,6 +19,7 @@ import json import tracemalloc from io import BytesIO +from types import SimpleNamespace from unittest.mock import patch import numpy as np @@ -41,8 +42,11 @@ from pypaimon.benchmark.act.experiment import load_experiment from pypaimon.benchmark.act.compare import canonical_sha256, compare_results from pypaimon.benchmark.act.harness import build_window_plan, run_backend +from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset from pypaimon.benchmark.act.paimon import ( + ACTION_COLUMNS, IMAGE_COLUMNS, + QPOS_COLUMNS, create_datasets as create_paimon_datasets, latest_snapshot_id, ) @@ -382,6 +386,92 @@ def test_independent_backend_results_preserve_tensor_and_loss_parity( assert comparison["experiments"][0]["backends"] == ["hdf5", "paimon"] +def test_backends_match_the_golden_act_window_contract(benchmark_input): + input_root, warehouse = benchmark_input + normalization = { + "qpos_mean": np.zeros(14, dtype=np.float32), + "qpos_std": np.ones(14, dtype=np.float32), + "action_mean": np.zeros(14, dtype=np.float32), + "action_std": np.ones(14, dtype=np.float32), + } + hdf5 = Hdf5ACTWindowDataset( + SimpleNamespace( + path=(input_root / "13_packbowl" / "success_episodes" / "train" + / "train-a" / "data" / "trajectory.hdf5"), + episode_id="train-a", + frame_count=6, + ), + normalization, + action_horizon=3, + ) + connection = pmm.connect( + database=agilex.DEFAULT_DATABASE, + options={"warehouse": str(warehouse)}, + ) + frames = connection.get_table(agilex.FRAMES_TABLE) + paimon, _ = create_paimon_datasets( + frames, + latest_snapshot_id(frames), + "train-a", + "val-a", + normalization, + BenchmarkConfig(action_horizon=3), + ) + + expected = hdf5[1] + actual = paimon[1] + + assert set(expected) == { + "sample_id", "episode_id", "step_idx", "qpos", "action", + "images", "is_pad", + } + assert expected["sample_id"] == "train-a#1" + assert expected["episode_id"] == "train-a" + assert expected["step_idx"] == 1 + assert torch.equal(expected["qpos"], torch.tensor( + list(range(408, 415)) + list(range(508, 515)), + dtype=torch.float32, + )) + assert torch.equal(expected["action"], torch.tensor([ + list(range(1208, 1215)) + list(range(1308, 1315)), + list(range(1215, 1222)) + list(range(1315, 1322)), + list(range(1222, 1229)) + list(range(1322, 1329)), + ], dtype=torch.float32)) + assert torch.allclose( + expected["images"][:, :, 0, 0], + torch.tensor([[2 / 255] * 3, [3 / 255] * 3, [4 / 255] * 3]), + ) + assert not expected["is_pad"].any() + for name in ("qpos", "action", "images", "is_pad"): + assert torch.equal(expected[name], actual[name]) + for name in ("sample_id", "episode_id", "step_idx"): + assert expected[name] == actual[name] + + +def test_hdf5_window_index_bounds(tmp_path): + path = _write_episode(tmp_path, "train", "train-a", 1) + dataset = Hdf5ACTWindowDataset( + SimpleNamespace( + path=path, + episode_id="train-a", + frame_count=6, + ), + { + "qpos_mean": np.zeros(14, dtype=np.float32), + "qpos_std": np.ones(14, dtype=np.float32), + "action_mean": np.zeros(14, dtype=np.float32), + "action_std": np.ones(14, dtype=np.float32), + }, + action_horizon=3, + ) + + assert dataset[-1]["sample_id"] == "train-a#3" + with pytest.raises(IndexError): + dataset[-len(dataset) - 1] + with pytest.raises(IndexError): + dataset[len(dataset)] + + def test_paimon_run_rejects_normalization_not_recorded_in_statistics( benchmark_input, tmp_path): input_root, warehouse = benchmark_input @@ -513,7 +603,14 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( assert fetch.call_count == 0 assert isinstance(train, ContiguousWindowDataset) assert isinstance(validation, ContiguousWindowDataset) - sample_before_append = train[0] + with patch.object( + train, "_read_rows", wraps=train._read_rows) as read_rows: + sample_before_append = train[0] + assert [call.args[1] for call in read_rows.call_args_list] == [ + list(ACTION_COLUMNS), + list(QPOS_COLUMNS + IMAGE_COLUMNS), + ] + assert [len(call.args[0]) for call in read_rows.call_args_list] == [3, 1] assert fetch.call_count == 1 assert { name: len(fetch.call_args.args[1][name]) From f65afe5a6ac92a46f76cccbea979eec94cd17baf Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 00:44:34 +0800 Subject: [PATCH 10/16] docs(python): clarify window dataset contracts Document contiguous ordering, anchor-column cardinality, batching, snapshot pinning, and map-style index behavior. Strengthen plural-access and pickle round-trip coverage. Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex AI-Contributed/Feature: 87/87 AI-Contributed/UT: 56/56 --- paimon-python/pypaimon/benchmark/act/hdf5.py | 5 ++ paimon-python/pypaimon/multimodal/query.py | 29 ++++++++-- .../pypaimon/multimodal/window_dataset.py | 53 ++++++++++++++++-- .../tests/contiguous_window_dataset_test.py | 56 ++++++++++++++----- 4 files changed, 118 insertions(+), 25 deletions(-) diff --git a/paimon-python/pypaimon/benchmark/act/hdf5.py b/paimon-python/pypaimon/benchmark/act/hdf5.py index cd46a733bad2..472ead7b4b9d 100644 --- a/paimon-python/pypaimon/benchmark/act/hdf5.py +++ b/paimon-python/pypaimon/benchmark/act/hdf5.py @@ -62,6 +62,11 @@ def __len__(self): return self.window_count def __getitem__(self, anchor): + """Return the ACT window whose first frame is ``anchor``. + + Negative anchors follow Python sequence semantics. State and images + come from the anchor frame, while action contains the complete horizon. + """ if anchor < 0: anchor += self.window_count if anchor < 0 or anchor >= self.window_count: diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 00bcb0f4fc31..33f0483b519a 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -182,10 +182,31 @@ def to_contiguous_window_dataset( The Dataset indexes only ``group_key``, ``order_key``, and Paimon row 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. + ``anchor_columns`` are provided to ``column_transforms`` as one-element + lists read from the first row of each window; ``adapter`` receives the + transformed values. ``order_key`` must contain non-null integers that + increase by exactly one within each group. The Dataset sorts rows within + each group and never creates a window across groups. + + Args: + window_size: Number of rows in a complete window. + columns: Value columns to return, excluding the group and order + keys. The scan projection is used when omitted. + anchor_columns: Subset of ``columns`` read only from the window's + first row. + group_key: Column identifying an independent row sequence. + order_key: Integer position column within each group. + stride: Distance between scheduled window starts. + tail: Handling for incomplete final windows: ``drop``, ``pad``, or + ``error``. + column_transforms: Per-column callables applied to value lists. + pad_values: Optional replacement values used by ``tail='pad'``. + adapter: Callable that converts the complete sample mapping. + blob_parallelism: Maximum concurrent BLOB body reads per fetch. + + Returns: + A snapshot-pinned ``ContiguousWindowDataset``. See that class for + padding, mask, transform, and adapter result semantics. """ if self._result_factory is not None: raise TypeError( diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index ed9da1f7a922..5fb5ccc5d4c1 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -37,17 +37,22 @@ class ContiguousWindowDataset(Dataset): 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``: + resolved while the index was built. Within each group, ``order_key`` must + contain non-null integers that increase by exactly one; rows from different + groups never share a window. ``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. - ``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. + The raw result mapping contains scalar group and order values, a + length-``window_size`` Boolean ``is_pad`` tensor, one-element lists for + ``anchor_columns``, and length-``window_size`` lists for other projected + columns. ``anchor_columns`` therefore avoids loading repeated context such + as observation images or initial robot state. ``column_transforms`` then + convert individual column lists before ``adapter`` adapts the complete + mapping to a model-specific contract. ``blob_parallelism`` controls concurrent BLOB reads for each item or batch. """ @@ -120,6 +125,12 @@ def __len__(self): return len(self._anchors) def __getitem__(self, index): + """Read one window by map-style Dataset index. + + Negative indices follow Python sequence semantics. The return value is + the pre-adapter mapping described by the class, or the adapter result + when an adapter is configured. + """ anchor, row_ids = self._resolve_window(index) rows = self._read_window_rows(row_ids) anchor_row = ( @@ -129,6 +140,11 @@ def __getitem__(self, index): return self._sample(anchor, rows, anchor_row) def __getitems__(self, indices): + """Read several Dataset indices while coalescing overlapping row IDs. + + The returned list preserves the requested index order and duplicates. + Coalescing affects only physical reads, not logical sample cardinality. + """ windows = [self._resolve_window(index) for index in indices] if not windows: return [] @@ -196,6 +212,17 @@ def _sample(self, anchor, rows, anchor_row=None): return sample def _build_index(self, index): + """Validate index rows and return grouped row IDs plus window anchors. + + Args: + index: Arrow table containing ``group_key``, ``order_key``, and + Paimon's ``_ROW_ID`` for the resolved snapshot. + + Returns: + ``(groups, anchors)``. Each group stores its key, ordered positions, + and row IDs. Each anchor stores group index, start offset, and the + number of real rows available before optional padding. + """ 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() @@ -264,6 +291,19 @@ def _read_window_rows(self, row_ids): return self._read_rows(row_ids, self._window_columns) def _read_rows(self, row_ids, columns=None): + """Read projected rows by ID from the pinned snapshot. + + Args: + row_ids: Paimon row IDs to read. Their order and duplicates define + the returned row order. + columns: Projected value columns, or all Dataset columns when + omitted. + + Returns: + A list of row dictionaries aligned one-for-one with ``row_ids``. + The internal ``_ROW_ID`` field is removed, and BLOB descriptors are + resolved to their bodies. + """ columns = self.columns if columns is None else columns query = ScanQuery(self._table) predicate_builder = ( @@ -324,6 +364,7 @@ def _read_window_index(query, group_key, order_key): def _pin_table(table, snapshot_id): + """Pin a table copy to ``snapshot_id``, or reuse it when unresolved.""" if snapshot_id is None: return table options = { diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index 0dc89e8d81fc..de332fcd232c 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -16,6 +16,7 @@ # under the License. import os +import pickle import shutil import tempfile import unittest @@ -159,24 +160,30 @@ def test_anchor_columns_read_only_the_window_anchor(self): 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]] + dataset = self._dataset( + self._table(), anchor_columns=["payload"]) 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"])) + actual = dataset.__getitems__([1, 0, 1]) + + self.assertEqual(2, read.call_count) + self.assertEqual(4, len(read.call_args_list[0].args[0])) + self.assertEqual(["value"], read.call_args_list[0].args[1]) + self.assertEqual(2, len(read.call_args_list[1].args[0])) + self.assertEqual(["payload"], read.call_args_list[1].args[1]) + self.assertEqual( + [("episode-b", 1), ("episode-b", 0), ("episode-b", 1)], + [(sample["episode"], sample["step"]) for sample in actual], + ) + self.assertEqual( + [[101, 102, 103], [100, 101, 102], [101, 102, 103]], + [sample["value"] for sample in actual], + ) + self.assertEqual( + [[b"episode-b-1"], [b"episode-b-0"], [b"episode-b-1"]], + [sample["payload"] for sample in actual], + ) def test_pad_tail_repeats_last_row_and_marks_real_padding(self): dataset = self._dataset( @@ -247,6 +254,25 @@ def test_pins_snapshot_for_later_on_demand_reads(self): self.assertEqual(2, len(dataset)) self.assertEqual([101, 102, 103], dataset[-1]["value"]) + def test_pickle_round_trip_preserves_snapshot_and_window(self): + dataset = self._dataset( + self._table(), anchor_columns=["payload"]) + expected = dataset[-1] + + restored = pickle.loads(pickle.dumps(dataset)) + + self.assertEqual(dataset.snapshot_id, restored.snapshot_id) + self.assertEqual( + dataset.snapshot_id, + restored._table.options.scan_snapshot_id(), + ) + actual = restored[-1] + self.assertEqual(expected["episode"], actual["episode"]) + self.assertEqual(expected["step"], actual["step"]) + self.assertEqual(expected["value"], actual["value"]) + self.assertEqual(expected["payload"], actual["payload"]) + self.assertTrue(torch.equal(expected["is_pad"], actual["is_pad"])) + def test_projection_filter_transform_and_dataloader_workers(self): table = self._table() dataset = ( From ac91b1c5768478de7eb8608fcb3164181d9bd755 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 02:03:09 +0800 Subject: [PATCH 11/16] fix(python): satisfy ACT runner style checks Add the required blank line before the nested dataset factory and remove the extra module-level blank line. Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 2/2 AI-Contributed/UT: 0/0 --- paimon-python/pypaimon/benchmark/act/runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/paimon-python/pypaimon/benchmark/act/runner.py b/paimon-python/pypaimon/benchmark/act/runner.py index 226e0dbdac3c..aad85fe27abc 100644 --- a/paimon-python/pypaimon/benchmark/act/runner.py +++ b/paimon-python/pypaimon/benchmark/act/runner.py @@ -417,6 +417,7 @@ def _paimon_factory_from_experiment( or not np.array_equal(action_std, normalization["action_std"])): raise ValueError( "Paimon normalization differs from the ACT experiment.") + def factory(): return create_paimon_datasets( frames, @@ -475,7 +476,6 @@ def _tensor_fingerprint(datasets, plan): } - def _shared_normalization( episodes, connection, From 09d00bcdfba21f8d6021edf81c084f2a0c41db01 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 02:55:13 +0800 Subject: [PATCH 12/16] fix(python): guard optional ACT benchmark imports Skip only the ACT runtime ownership checks when Torch or Pillow is absent, while keeping dependency-free comparison tests active. Co-Authored-By: Codex AI-Model: gpt-5.6-sol Co-Authored-By: Codex Co-Authored-By: Codex AI-Contributed/Feature: 0/0 AI-Contributed/UT: 8/8 --- paimon-python/pypaimon/tests/act_benchmark_test.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/paimon-python/pypaimon/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py index 81d55bc1708c..d0261127b3a5 100644 --- a/paimon-python/pypaimon/tests/act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/act_benchmark_test.py @@ -185,12 +185,14 @@ def test_compare_rejects_tampered_result_experiment_hash(): def test_shared_harness_is_imported_from_act_package(): + _require_act_runtime() from pypaimon.benchmark.act.harness import BenchmarkConfig assert BenchmarkConfig().to_dict() == load_experiment()["config"] def test_hdf5_dataset_is_owned_by_hdf5_backend(): + _require_act_runtime() from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset assert Hdf5ACTWindowDataset.__module__ == ( @@ -198,12 +200,18 @@ def test_hdf5_dataset_is_owned_by_hdf5_backend(): def test_paimon_adapter_is_owned_by_paimon_backend(): + _require_act_runtime() from pypaimon.benchmark.act.paimon import PaimonACTAdapter assert PaimonACTAdapter.__module__ == ( "pypaimon.benchmark.act.paimon") +def _require_act_runtime(): + pytest.importorskip("torch") + pytest.importorskip("PIL.Image") + + def _result(backend, experiment, environment, throughput): return { "schema_version": "act-benchmark-result@1", From dfbf45944f79944ac9e18e74f7e7cab2d965a0fc Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 11:19:40 +0800 Subject: [PATCH 13/16] fix(python): stabilize ACT throughput measurement Measure four complete physical fetch groups by default so the first Paimon fetch does not dominate steady-state throughput. Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 4/4 AI-Contributed/UT: 9/9 --- .../pypaimon/benchmark/act/default_experiment.json | 2 +- paimon-python/pypaimon/benchmark/act/harness.py | 2 +- paimon-python/pypaimon/tests/act_benchmark_test.py | 9 +++++++++ 3 files changed, 11 insertions(+), 2 deletions(-) diff --git a/paimon-python/pypaimon/benchmark/act/default_experiment.json b/paimon-python/pypaimon/benchmark/act/default_experiment.json index df3241ca683d..91c2ffe291c0 100644 --- a/paimon-python/pypaimon/benchmark/act/default_experiment.json +++ b/paimon-python/pypaimon/benchmark/act/default_experiment.json @@ -10,7 +10,7 @@ "optimizer_steps": 2, "rounds": 3, "seed": 20260825, - "timed_batches": 4, + "timed_batches": 32, "warmup_batches": 1, "weight_decay": 0.0001 }, diff --git a/paimon-python/pypaimon/benchmark/act/harness.py b/paimon-python/pypaimon/benchmark/act/harness.py index fa8164765a4f..d9b41539f48c 100644 --- a/paimon-python/pypaimon/benchmark/act/harness.py +++ b/paimon-python/pypaimon/benchmark/act/harness.py @@ -58,7 +58,7 @@ class BenchmarkConfig: learning_rate: float = 1e-4 weight_decay: float = 1e-4 warmup_batches: int = 1 - timed_batches: int = 4 + timed_batches: int = 32 fetch_batches: int = 8 rounds: int = 3 diff --git a/paimon-python/pypaimon/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py index d0261127b3a5..31d9c3dde3c0 100644 --- a/paimon-python/pypaimon/tests/act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/act_benchmark_test.py @@ -33,11 +33,20 @@ def test_default_experiment_loads_packaged_benchmark_parameters(): assert experiment["benchmark_id"] == "robomind-act" assert experiment["config"]["seed"] == 20260825 assert experiment["config"]["action_horizon"] == 32 + assert experiment["config"]["timed_batches"] == 32 assert experiment["config"]["rounds"] == 3 assert experiment["statistics_version"] == ( "robomind-agilex-joint-position@1") +def test_default_throughput_measurement_spans_four_physical_fetches(): + """Keep the first Paimon fetch from dominating steady-state throughput.""" + config = load_experiment()["config"] + + assert config["timed_batches"] % config["fetch_batches"] == 0 + assert config["timed_batches"] // config["fetch_batches"] == 4 + + def test_compare_reports_ratio_for_matching_backend_results(): experiment = {"schema_version": "act-benchmark-experiment@1"} environment = {"python": "3.10", "machine": "arm64"} From 06fd1c279cbe732ece040d860e1a70c507c2e793 Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 11:44:49 +0800 Subject: [PATCH 14/16] test(python): simplify ACT benchmark coverage Avoid coupling tests to tuned experiment values and fold redundant integration assertions into the paired backend parity test. Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 0/0 AI-Contributed/UT: 174/174 --- .../pypaimon/tests/act_benchmark_test.py | 80 ++++------------ .../pypaimon/tests/act_runner_test.py | 94 ++----------------- 2 files changed, 28 insertions(+), 146 deletions(-) diff --git a/paimon-python/pypaimon/tests/act_benchmark_test.py b/paimon-python/pypaimon/tests/act_benchmark_test.py index 31d9c3dde3c0..9f9fb4357c09 100644 --- a/paimon-python/pypaimon/tests/act_benchmark_test.py +++ b/paimon-python/pypaimon/tests/act_benchmark_test.py @@ -26,36 +26,31 @@ from pypaimon.benchmark.act.experiment import load_experiment -def test_default_experiment_loads_packaged_benchmark_parameters(): +def test_packaged_experiment_contains_a_valid_benchmark_config(): + _require_act_runtime() + from pypaimon.benchmark.act.harness import BenchmarkConfig + experiment = load_experiment() assert experiment["schema_version"] == "act-benchmark-experiment@1" assert experiment["benchmark_id"] == "robomind-act" - assert experiment["config"]["seed"] == 20260825 - assert experiment["config"]["action_horizon"] == 32 - assert experiment["config"]["timed_batches"] == 32 - assert experiment["config"]["rounds"] == 3 + assert BenchmarkConfig(**experiment["config"]).to_dict() == ( + experiment["config"]) assert experiment["statistics_version"] == ( "robomind-agilex-joint-position@1") -def test_default_throughput_measurement_spans_four_physical_fetches(): - """Keep the first Paimon fetch from dominating steady-state throughput.""" - config = load_experiment()["config"] - - assert config["timed_batches"] % config["fetch_batches"] == 0 - assert config["timed_batches"] // config["fetch_batches"] == 4 - - def test_compare_reports_ratio_for_matching_backend_results(): experiment = {"schema_version": "act-benchmark-experiment@1"} environment = {"python": "3.10", "machine": "arm64"} - results = [ - _result("hdf5", experiment, environment, throughput=10.0), - _result("paimon", experiment, environment, throughput=15.0), - ] + hdf5 = _result("hdf5", experiment, environment, throughput=10.0) + paimon = _result("paimon", experiment, environment, throughput=15.0) + hdf5["summary"]["first_batch_s"] = { + "median": 2.0, "min": 2.0, "max": 2.0} + paimon["summary"]["first_batch_s"] = { + "median": 1.0, "min": 1.0, "max": 1.0} - comparison = compare_results(results) + comparison = compare_results([hdf5, paimon]) assert comparison["status"] == "SUCCEEDED" assert len(comparison["experiments"]) == 1 @@ -67,6 +62,12 @@ def test_compare_reports_ratio_for_matching_backend_results(): "paimon_over_hdf5": 1.5, "preferred": "higher", } + assert group["metrics"]["first_batch_s"] == { + "hdf5": 2.0, + "paimon": 1.0, + "hdf5_over_paimon": 2.0, + "preferred": "lower", + } def test_compare_rejects_different_tensor_fingerprints(): @@ -100,26 +101,6 @@ def test_compare_requires_results_from_both_backends(): assert group["metrics"] == {} -def test_compare_reports_lower_is_better_metric_as_paimon_speedup(): - experiment = {"schema_version": "act-benchmark-experiment@1"} - environment = {"python": "3.10", "machine": "arm64"} - hdf5 = _result("hdf5", experiment, environment, throughput=10.0) - paimon = _result("paimon", experiment, environment, throughput=15.0) - hdf5["summary"]["first_batch_s"] = { - "median": 2.0, "min": 2.0, "max": 2.0} - paimon["summary"]["first_batch_s"] = { - "median": 1.0, "min": 1.0, "max": 1.0} - - comparison = compare_results([hdf5, paimon]) - - assert comparison["experiments"][0]["metrics"]["first_batch_s"] == { - "hdf5": 2.0, - "paimon": 1.0, - "hdf5_over_paimon": 2.0, - "preferred": "lower", - } - - def test_load_results_combines_explicit_files_and_directory(tmp_path): experiment = {"schema_version": "act-benchmark-experiment@1"} environment = {"python": "3.10", "machine": "arm64"} @@ -193,29 +174,6 @@ def test_compare_rejects_tampered_result_experiment_hash(): compare_results([result]) -def test_shared_harness_is_imported_from_act_package(): - _require_act_runtime() - from pypaimon.benchmark.act.harness import BenchmarkConfig - - assert BenchmarkConfig().to_dict() == load_experiment()["config"] - - -def test_hdf5_dataset_is_owned_by_hdf5_backend(): - _require_act_runtime() - from pypaimon.benchmark.act.hdf5 import Hdf5ACTWindowDataset - - assert Hdf5ACTWindowDataset.__module__ == ( - "pypaimon.benchmark.act.hdf5") - - -def test_paimon_adapter_is_owned_by_paimon_backend(): - _require_act_runtime() - from pypaimon.benchmark.act.paimon import PaimonACTAdapter - - assert PaimonACTAdapter.__module__ == ( - "pypaimon.benchmark.act.paimon") - - def _require_act_runtime(): pytest.importorskip("torch") pytest.importorskip("PIL.Image") diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py index 67e1d88fe9ed..d1bb02928d50 100644 --- a/paimon-python/pypaimon/tests/act_runner_test.py +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -55,10 +55,6 @@ from pypaimon.sample import robomind_agilex as agilex -def test_default_fetch_group_covers_eight_logical_batches(): - assert BenchmarkConfig().fetch_batches == 8 - - def test_logical_batches_coalesce_one_physical_fetch(): class BatchDataset: def __init__(self): @@ -302,44 +298,6 @@ def test_prepare_writes_resolved_experiment(benchmark_input, tmp_path): assert experiment["paimon"]["frames_snapshot_id"] > 0 -def test_hdf5_run_consumes_resolved_experiment_without_warehouse( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "seed": 17, - "action_horizon": 3, - "batch_size": 2, - "optimizer_steps": 2, - "image_height": 8, - "image_width": 10, - "timed_batches": 2, - }) - experiment_path = tmp_path / "experiment.json" - experiment = prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - result_path = tmp_path / "hdf5-result.json" - - result = run_experiment( - "hdf5", - experiment_path, - result_path, - input_root=input_root, - policy_factory=_policy_factory, - ) - - assert json.loads(result_path.read_text()) == result - assert result["schema_version"] == "act-benchmark-result@1" - assert result["status"] == "SUCCEEDED" - assert result["backend"] == "hdf5" - assert result["experiment"] == experiment - assert len(result["experiment_sha256"]) == 64 - assert len(result["tensor_fingerprint"]["sha256"]) == 64 - assert len(result["runs"]) == 3 - assert result["summary"]["round_count"] == 3 - - def test_independent_backend_results_preserve_tensor_and_loss_parity( benchmark_input, tmp_path): input_root, warehouse = benchmark_input @@ -358,21 +316,28 @@ def test_independent_backend_results_preserve_tensor_and_loss_parity( prepare_experiment( input_root, warehouse, experiment_path, definition=definition) + hdf5_path = tmp_path / "hdf5-result.json" hdf5_result = run_experiment( "hdf5", experiment_path, - tmp_path / "hdf5-result.json", + hdf5_path, input_root=input_root, policy_factory=_policy_factory, ) + paimon_path = tmp_path / "paimon-result.json" paimon_result = run_experiment( "paimon", experiment_path, - tmp_path / "paimon-result.json", + paimon_path, warehouse=warehouse, policy_factory=_policy_factory, ) + assert json.loads(hdf5_path.read_text()) == hdf5_result + assert json.loads(paimon_path.read_text()) == paimon_result + assert hdf5_result["schema_version"] == "act-benchmark-result@1" + assert paimon_result["schema_version"] == "act-benchmark-result@1" + assert hdf5_result["experiment"] == paimon_result["experiment"] assert hdf5_result["tensor_fingerprint"] == ( paimon_result["tensor_fingerprint"]) assert [run["train_loss"] for run in hdf5_result["runs"]] == [ @@ -636,47 +601,6 @@ def test_paimon_windows_are_lazy_snapshot_pinned_and_vortex_independent( sample_before_append[name], sample_after_append[name]) -def test_compare_rejects_different_hdf5_tensor_bytes( - benchmark_input, tmp_path): - input_root, warehouse = benchmark_input - definition = load_experiment() - definition["statistics_version"] = "act-test@1" - definition["config"].update({ - "action_horizon": 3, - "batch_size": 1, - "optimizer_steps": 1, - "image_height": 8, - "image_width": 10, - }) - experiment_path = tmp_path / "experiment.json" - prepare_experiment( - input_root, warehouse, experiment_path, definition=definition) - paimon_result = run_experiment( - "paimon", - experiment_path, - tmp_path / "paimon-result.json", - warehouse=warehouse, - policy_factory=_policy_factory, - ) - 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 - - hdf5_result = run_experiment( - "hdf5", - experiment_path, - tmp_path / "hdf5-result.json", - input_root=input_root, - policy_factory=_policy_factory, - ) - - comparison = compare_results([hdf5_result, paimon_result]) - assert comparison["status"] == "FAILED" - assert comparison["experiments"][0]["reason"] == ( - "tensor fingerprints differ") - - def test_requires_at_least_three_measurement_rounds(): with pytest.raises(ValueError, match="rounds must be at least 3"): BenchmarkConfig(rounds=2) From c1dbea125f95d741637a3618e18266e33b800f7a Mon Sep 17 00:00:00 2001 From: Yann Date: Wed, 2 Sep 2026 17:14:14 +0800 Subject: [PATCH 15/16] fix(python): address ACT benchmark review feedback Resolve deferred MAP BLOB reads, snapshot option conflicts, and shared mutable window cells. Tighten fetch timing, runtime identity, and Python version reporting. Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 100/100 AI-Contributed/UT: 149/149 --- docs/docs/pypaimon/robomind-act-benchmark.md | 2 + .../pypaimon/benchmark/act/__main__.py | 13 +++ .../pypaimon/benchmark/act/harness.py | 5 +- .../pypaimon/benchmark/act/runner.py | 62 ++++++++++-- .../pypaimon/multimodal/window_dataset.py | 18 +++- .../pypaimon/tests/act_runner_test.py | 50 ++++++++-- .../tests/contiguous_window_dataset_test.py | 99 +++++++++++++++++++ 7 files changed, 227 insertions(+), 22 deletions(-) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index bb40a8c4277f..fdb9a460184c 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -36,6 +36,8 @@ calculates performance ratios. ## Install +Python 3.10 or newer is required. + ```shell pip install 'pypaimon[act,hdf5]' ``` diff --git a/paimon-python/pypaimon/benchmark/act/__main__.py b/paimon-python/pypaimon/benchmark/act/__main__.py index 686708e42072..d1d8560ae275 100644 --- a/paimon-python/pypaimon/benchmark/act/__main__.py +++ b/paimon-python/pypaimon/benchmark/act/__main__.py @@ -16,6 +16,19 @@ """Command-line entry point for ACT benchmark preparation, runs, and reports.""" +# ruff: noqa: E402 + +import sys + + +def _require_supported_python(version_info): + """Reject runtimes older than the ACT dependencies support.""" + if tuple(version_info[:2]) < (3, 10): + raise RuntimeError("ACT benchmark requires Python 3.10 or newer.") + + +_require_supported_python(sys.version_info) + import argparse import copy import json diff --git a/paimon-python/pypaimon/benchmark/act/harness.py b/paimon-python/pypaimon/benchmark/act/harness.py index d9b41539f48c..8fe7fce273dc 100644 --- a/paimon-python/pypaimon/benchmark/act/harness.py +++ b/paimon-python/pypaimon/benchmark/act/harness.py @@ -381,13 +381,14 @@ def run_backend( logical_batch_size=config.batch_size, fetch_batches=config.fetch_batches, ) - batch_fetch_started = time.monotonic() + batch_fetch_seconds = 0.0 batch_fetch_sample_count = 0 for _ in range(config.timed_batches): + batch_fetch_started = time.monotonic() batch = next(batch_fetch_iterator) + batch_fetch_seconds += time.monotonic() - batch_fetch_started validate_act_batch(batch, config) batch_fetch_sample_count += len(batch["sample_id"]) - batch_fetch_seconds = time.monotonic() - batch_fetch_started _seed_everything(config.seed) policy, model = policy_factory(config) diff --git a/paimon-python/pypaimon/benchmark/act/runner.py b/paimon-python/pypaimon/benchmark/act/runner.py index aad85fe27abc..d82ca1683954 100644 --- a/paimon-python/pypaimon/benchmark/act/runner.py +++ b/paimon-python/pypaimon/benchmark/act/runner.py @@ -37,9 +37,13 @@ from datetime import datetime, timezone from pathlib import Path +import PIL +import h5py import numpy as np +import pyarrow as pa import torch import pypaimon.multimodal as pmm +from pypaimon import build_info from pypaimon.benchmark.act.hdf5 import ( compute_normalization as compute_hdf5_normalization, create_datasets as create_hdf5_datasets, @@ -264,13 +268,8 @@ def dataset_factory(): "model": runs[0]["model"], "runs": runs, "summary": _summarize(runs), - "environment": { - "python": platform.python_version(), - "os": platform.platform(), - "machine": platform.machine(), - "torch": torch.__version__, - "source_commit": _git_head(Path(__file__).resolve().parents[4]), - }, + "environment": _runtime_environment( + Path(__file__).resolve().parents[4]), "command": _command_argv(), "timing": {"wall_time_s": time.monotonic() - started}, "unverified": [ @@ -680,6 +679,55 @@ def _git_head(repository): return "UNKNOWN" +def _runtime_environment(repository): + """Return dependency, CPU, thread, and source identity for comparison.""" + source_commit = _git_head(repository) + package_build = build_info.full_version() + if source_commit == "UNKNOWN" and package_build == "UNKNOWN": + raise RuntimeError( + "ACT benchmark cannot determine its source identity.") + return { + "python": platform.python_version(), + "os": platform.platform(), + "machine": platform.machine(), + "cpu_identity": _cpu_identity(), + "cpu_count": os.cpu_count() or 1, + "torch_threads": torch.get_num_threads(), + "torch_interop_threads": torch.get_num_interop_threads(), + "pypaimon_build": package_build, + "numpy": np.__version__, + "pyarrow": pa.__version__, + "h5py": h5py.__version__, + "pillow": PIL.__version__, + "torch": torch.__version__, + "source_commit": source_commit, + } + + +def _cpu_identity(): + """Return the most specific CPU model available from the local OS.""" + if platform.system() == "Darwin": + try: + return subprocess.check_output( + ["sysctl", "-n", "machdep.cpu.brand_string"], + stderr=subprocess.DEVNULL, + universal_newlines=True, + ).strip() + except (OSError, subprocess.CalledProcessError): + pass + identity = platform.processor().strip() + if identity: + return identity + if platform.system() == "Linux": + try: + for line in Path("/proc/cpuinfo").read_text().splitlines(): + if line.startswith(("model name", "Hardware")): + return line.partition(":")[2].strip() + except OSError: + pass + return platform.machine() + + def _command_argv(): """Return the invoked Python basename and command-line arguments.""" import sys diff --git a/paimon-python/pypaimon/multimodal/window_dataset.py b/paimon-python/pypaimon/multimodal/window_dataset.py index 5fb5ccc5d4c1..6eab515a2cbb 100644 --- a/paimon-python/pypaimon/multimodal/window_dataset.py +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -27,7 +27,7 @@ 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.schema.data_types import is_blob_type, is_map_blob_type from pypaimon.snapshot.time_travel_util import SCAN_KEYS from pypaimon.table.special_fields import SpecialFields @@ -198,9 +198,9 @@ def _sample(self, anchor, rows, anchor_row=None): } for name in self.columns: if name in self.anchor_columns: - values = [anchor_row[name]] + values = [copy.deepcopy(anchor_row[name])] else: - values = [row[name] for row in rows] + values = [copy.deepcopy(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( @@ -320,7 +320,8 @@ def _read_rows(self, row_ids, columns=None): blob_columns = [ field.name for field in self._table.fields - if field.name in columns and is_blob_type(field.type) + if field.name in columns + and (is_blob_type(field.type) or is_map_blob_type(field.type)) ] if blob_columns: scalar, blobs = query.read_blobs( @@ -367,8 +368,15 @@ def _pin_table(table, snapshot_id): """Pin a table copy to ``snapshot_id``, or reuse it when unresolved.""" if snapshot_id is None: return table + scan_keys = set(SCAN_KEYS) + scan_keys.update(option.key() for option in ( + CoreOptions.SCAN_MODE, + CoreOptions.INCREMENTAL_BETWEEN_TIMESTAMP, + CoreOptions.SCAN_FILE_CREATION_TIME_MILLIS, + CoreOptions.SCAN_CREATION_TIME_MILLIS, + )) options = { - key: None for key in SCAN_KEYS + key: None for key in scan_keys if table.options.options.contains_key(key) } options[CoreOptions.SCAN_SNAPSHOT_ID.key()] = str(snapshot_id) diff --git a/paimon-python/pypaimon/tests/act_runner_test.py b/paimon-python/pypaimon/tests/act_runner_test.py index d1bb02928d50..e8f14c674df6 100644 --- a/paimon-python/pypaimon/tests/act_runner_test.py +++ b/paimon-python/pypaimon/tests/act_runner_test.py @@ -33,9 +33,9 @@ import pypaimon.multimodal as pmm import pypaimon.benchmark.act.harness as act_harness import pypaimon.benchmark.act.__main__ as act_cli +import pypaimon.benchmark.act.runner as act_runner from pypaimon.benchmark.act.runner import ( BenchmarkConfig, - _git_head, prepare_experiment, run_experiment, ) @@ -155,6 +155,8 @@ def test_backend_coalesces_timed_batch_fetches(): rounds=3, ) + clock = SimpleNamespace(value=0.0) + class BatchDataset(torch.utils.data.Dataset): def __init__(self): self.calls = [] @@ -174,14 +176,23 @@ def __getitem__(self, index): } def __getitems__(self, indices): + clock.value += 0.25 self.calls.append(list(indices)) return [self[index] for index in indices] dataset = BatchDataset() plan = build_window_plan(len(dataset), len(dataset), config) - with patch.object(act_harness, "_measure_python_peak", return_value=0): - run_backend( + def validate(_batch, _config): + clock.value += 1.0 + + with ( + patch.object(act_harness, "_measure_python_peak", return_value=0), + patch.object( + act_harness.time, "monotonic", side_effect=lambda: clock.value), + patch.object(act_harness, "validate_act_batch", side_effect=validate), + ): + result = run_backend( "test", 1, lambda: (dataset, dataset), @@ -197,6 +208,7 @@ def __getitems__(self, indices): list(plan.train_indices), list(plan.validation_indices), ] + assert result["batch_fetch_s"] == 0.25 def _jpeg(value): @@ -634,6 +646,13 @@ def test_cli_exposes_prepare_run_and_compare_contracts(capsys): assert "--results-dir" in capsys.readouterr().out +def test_cli_requires_python_3_10_or_newer(): + with pytest.raises(RuntimeError, match="Python 3.10 or newer"): + act_cli._require_supported_python((3, 9)) + + act_cli._require_supported_python((3, 10)) + + def test_automatic_artifact_paths_do_not_overwrite_same_second(tmp_path): with patch.object(act_cli, "datetime") as now: now.now.return_value.strftime.return_value = "20260901T120000Z" @@ -646,8 +665,23 @@ def test_automatic_artifact_paths_do_not_overwrite_same_second(tmp_path): assert second.parent == tmp_path -def test_source_commit_falls_back_outside_git_checkout(tmp_path): - with patch( - "pypaimon.benchmark.act.runner.subprocess.check_output", - side_effect=FileNotFoundError): - assert _git_head(tmp_path) == "UNKNOWN" +def test_runtime_environment_uses_package_identity_outside_git_checkout( + tmp_path): + with patch.object(act_runner, "_git_head", return_value="UNKNOWN"): + environment = act_runner._runtime_environment(tmp_path) + + assert environment["source_commit"] == "UNKNOWN" + assert environment["pypaimon_build"] != "UNKNOWN" + assert environment["cpu_identity"] + assert environment["cpu_count"] > 0 + assert environment["torch_threads"] > 0 + assert environment["torch_interop_threads"] > 0 + assert all(environment[name] for name in ( + "numpy", "pyarrow", "h5py", "pillow")) + + with ( + patch.object(act_runner, "_git_head", return_value="UNKNOWN"), + patch.object(act_runner.build_info, "full_version", return_value="UNKNOWN"), + pytest.raises(RuntimeError, match="source identity"), + ): + act_runner._runtime_environment(tmp_path) diff --git a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py index de332fcd232c..389d4bb5f36a 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -146,6 +146,56 @@ def test_reads_blob_payloads_only_when_a_window_is_requested(self): sample["payload"], ) + def test_reads_map_blob_payloads_for_map_only_and_mixed_windows(self): + schema = pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("payload", pa.large_binary()), + pa.field( + "attachments", + pa.map_(pa.string(), pa.large_binary()), + ), + ]) + table = self.conn.create_table( + "map_blobs", + schema=schema, + options=_TABLE_OPTIONS, + ) + table.add(pa.Table.from_pylist([ + { + "episode": "episode-a", + "step": 0, + "payload": b"scalar-0", + "attachments": { + "body": b"map-0", "empty": b"", "null": None}, + }, + { + "episode": "episode-a", + "step": 1, + "payload": b"scalar-1", + "attachments": None, + }, + ], schema=schema)) + + def window(columns): + return table.scan().to_contiguous_window_dataset( + window_size=2, + columns=columns, + group_key="episode", + order_key="step", + )[0] + + map_only = window(["attachments"]) + mixed = window(["payload", "attachments"]) + + self.assertEqual( + {"body": b"map-0", "empty": b"", "null": None}, + dict(map_only["attachments"][0]), + ) + self.assertIsNone(map_only["attachments"][1]) + self.assertEqual([b"scalar-0", b"scalar-1"], mixed["payload"]) + self.assertEqual(map_only["attachments"], mixed["attachments"]) + def test_anchor_columns_read_only_the_window_anchor(self): table = self._table() original = ScanQuery._fetch_bodies @@ -185,6 +235,42 @@ def test_plural_access_coalesces_overlapping_window_reads(self): [sample["payload"] for sample in actual], ) + def test_plural_access_isolates_mutable_cells_between_samples(self): + table = self.conn.create_table( + "mutable_cells", + schema=pa.schema([ + pa.field("episode", pa.string(), nullable=False), + pa.field("step", pa.int32(), nullable=False), + pa.field("values", pa.list_(pa.int32()), nullable=False), + ]), + options=_TABLE_OPTIONS, + ) + table.add([ + {"episode": "episode-a", "step": step, "values": [step]} + for step in range(3) + ]) + + def mutate(values): + for value in values: + value.append(99) + return values + + dataset = table.scan().to_contiguous_window_dataset( + window_size=2, + columns=["values"], + group_key="episode", + order_key="step", + column_transforms={"values": mutate}, + ) + + batched = dataset.__getitems__([0, 1, 0]) + singles = [dataset[index] for index in (0, 1, 0)] + + self.assertEqual( + [sample["values"] for sample in singles], + [sample["values"] for sample in batched], + ) + def test_pad_tail_repeats_last_row_and_marks_real_padding(self): dataset = self._dataset( self._table(), tail="pad", pad_values={"value": -1}) @@ -254,6 +340,19 @@ def test_pins_snapshot_for_later_on_demand_reads(self): self.assertEqual(2, len(dataset)) self.assertEqual([101, 102, 103], dataset[-1]["value"]) + def test_snapshot_pin_clears_scan_mode_before_on_demand_reads(self): + table = self._table() + query = ScanQuery(table.raw_table.copy({"scan.mode": "latest-full"})) + + dataset = query.to_contiguous_window_dataset( + window_size=3, + columns=["value", "payload"], + group_key="episode", + order_key="step", + ) + + self.assertEqual([100, 101, 102], dataset[0]["value"]) + def test_pickle_round_trip_preserves_snapshot_and_window(self): dataset = self._dataset( self._table(), anchor_columns=["payload"]) From 9e0340a82267eec221872df479551d7f0619448e Mon Sep 17 00:00:00 2001 From: Yann Date: Thu, 3 Sep 2026 10:45:02 +0800 Subject: [PATCH 16/16] docs(python): clarify ACT benchmark measurements Document complete experiment overrides and distinguish end-to-end optimizer timing from per-step compute timing. Co-Authored-By: Codex AI-Model: gpt-5.6-sol AI-Contributed/Feature: 15/15 AI-Contributed/UT: 0/0 --- docs/docs/pypaimon/robomind-act-benchmark.md | 8 +++++--- paimon-python/pypaimon/benchmark/act/harness.py | 7 +++++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md index fdb9a460184c..37f9e011fbbc 100644 --- a/docs/docs/pypaimon/robomind-act-benchmark.md +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -58,8 +58,9 @@ snapshot, and materializes deterministic measurement, training, and validation window indices. Without `--experiment`, preparation starts from the packaged -`default_experiment.json`. A custom JSON file can change the defaults, and -individual values can be overridden on the command line: +`default_experiment.json`. `--experiment` replaces that definition, so a +custom JSON file must contain every required field. Command-line options then +override individual values: ```shell python -m pypaimon.benchmark.act prepare \ @@ -154,7 +155,8 @@ Every backend repeat records: - dataset construction time; - first-batch latency after construction; - batch-fetch samples per second after warm-up; -- fixed ACT optimizer-step time and per-step loss trace; +- end-to-end fixed ACT optimizer-step time, including dataset fetch; +- per-step loss and compute time after each training batch has been fetched; - validation loss; - total measured wall time; - Python peak allocation from a separate dataset-first-batch replay. diff --git a/paimon-python/pypaimon/benchmark/act/harness.py b/paimon-python/pypaimon/benchmark/act/harness.py index 8fe7fce273dc..ee1aca7a3772 100644 --- a/paimon-python/pypaimon/benchmark/act/harness.py +++ b/paimon-python/pypaimon/benchmark/act/harness.py @@ -351,8 +351,11 @@ def run_backend( Returns: A JSON-compatible metrics dictionary covering dataset construction, first batch, timed batch fetch, fixed optimizer steps, validation loss, - and a separate ``tracemalloc`` peak replay. OS page cache is not - controlled and native Arrow/Torch allocations are outside tracemalloc. + and a separate ``tracemalloc`` peak replay. ``fixed_steps_s`` includes + dataset fetch, while each ``train_trace.step_time_s`` starts after its + batch is fetched and covers conversion, forward/backward, and optimizer + update. OS page cache is not controlled and native Arrow/Torch + allocations are outside tracemalloc. """ _seed_everything(config.seed) policy_factory = policy_factory or build_act_policy