diff --git a/docs/docs/pypaimon/multimodal-api.mdx b/docs/docs/pypaimon/multimodal-api.mdx index e1074ee5ce8c..f173afc89fc8 100644 --- a/docs/docs/pypaimon/multimodal-api.mdx +++ b/docs/docs/pypaimon/multimodal-api.mdx @@ -720,6 +720,64 @@ Notes: few large reads); scattered point reads coalesce less. - Blob reads are available only on `scan()`, not on the `search()` queries. +### Contiguous windows for PyTorch + +Install the `torch` extra, then use `to_contiguous_window_dataset` to expose +map-style windows without loading the selected rows or BLOB payloads into Python +memory up front. The Dataset builds a compact index from the group column, order +column, and Paimon row IDs. Each `__getitem__` call fetches only that window from +the snapshot recorded in `dataset.snapshot_id`. + +```shell +pip install pypaimon[torch] +``` + +```python +import torch + + +def float32_window(values): + return torch.tensor(values, dtype=torch.float32) + + +windows = ( + frames.scan() + .where("split = 'train'") + .to_contiguous_window_dataset( + window_size=16, + columns=["state", "action"], + group_key="episode_id", + order_key="step_idx", + tail="pad", + column_transforms={ + "state": float32_window, + "action": float32_window, + }, + ) +) + +sample = windows[0] +assert sample["action"].shape == (16, action_size) +assert sample["is_pad"].shape == (16,) +``` + +The group and order keys in a sample identify the window anchor. Every projected +column contains the whole window. With `tail="drop"`, only full windows are +exposed. With `tail="pad"`, every real row is an anchor; missing suffix values +repeat the last real value by default and `is_pad` is `True` exactly at those +positions. With `tail="error"`, construction fails if any scheduled anchor is +incomplete. Use `pad_values` to override the repeated value for individual +columns. Anchors advance by `stride`, which defaults to one row. + +`column_transforms` receive one padded Python list per projected column. This is +where applications define tensor dtype and shape or decode BLOB bytes. The +optional `adapter` receives the resulting sample mapping and can rename or +combine fields for a model-specific batch contract. The core Dataset does not +know model field names, image formats, or normalization rules. Top-level +functions and callable classes are recommended for transforms and adapters so +the Dataset remains picklable by multi-worker `torch.utils.data.DataLoader` +instances. + ### Distributed BLOB processing with Ray For larger jobs, read descriptors with `to_ray()`, then fetch and process BLOB diff --git a/docs/docs/pypaimon/pytorch.md b/docs/docs/pypaimon/pytorch.md index ac9f6b8d64eb..26020d72fe09 100644 --- a/docs/docs/pypaimon/pytorch.md +++ b/docs/docs/pypaimon/pytorch.md @@ -106,7 +106,49 @@ 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"], + anchor_columns=["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. 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. + +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/docs/docs/pypaimon/robomind-act-benchmark.md b/docs/docs/pypaimon/robomind-act-benchmark.md new file mode 100644 index 000000000000..664465fc111e --- /dev/null +++ b/docs/docs/pypaimon/robomind-act-benchmark.md @@ -0,0 +1,83 @@ +--- +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 +``` + +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 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 +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 +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..dd874406579c --- /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("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()): + 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 reduced CPU LeRobot ACT configuration used by the benchmark.""" + 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..93cfa512d05b --- /dev/null +++ b/paimon-python/pypaimon/benchmark/paired_act.py @@ -0,0 +1,865 @@ +# 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, + 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.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): + 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 DataLoader 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="DataLoader batches consumed before timing.") + parser.add_argument( + "--loader-batches", type=int, default=BenchmarkConfig.loader_batches, + help="Batches used for DataLoader throughput measurement.") + 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, + 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/query.py b/paimon-python/pypaimon/multimodal/query.py index 9046498f7cf8..310c0c1e11f0 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -151,6 +151,49 @@ def to_torch( max_buffer_input_splits=max_buffer_input_splits, ) + def to_contiguous_window_dataset( + self, + *, + window_size, + columns=None, + anchor_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. 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. + """ + 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, + anchor_columns=anchor_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..ed9da1f7a922 --- /dev/null +++ b/paimon-python/pypaimon/multimodal/window_dataset.py @@ -0,0 +1,442 @@ +# 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. + + ``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") + + def __init__( + self, + query, + *, + window_size, + columns=None, + anchor_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.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) + 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.") + + 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): + anchor, row_ids = self._resolve_window(index) + 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] + 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_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 + ] + + 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") + + 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, anchor_row=None): + 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: + 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: + 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)) + 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_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() + .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(columns) + query._include_row_id = True + + 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( + blob_columns, parallelism=self.blob_parallelism) + rows = scalar.to_pylist() + for name in 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_key, order_key): + index_query = copy.copy(query) + 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() + 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 _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") + 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..0dc89e8d81fc --- /dev/null +++ b/paimon-python/pypaimon/tests/contiguous_window_dataset_test.py @@ -0,0 +1,347 @@ +# 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_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]] + + 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}) + + 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() 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..ca49d3068ef5 --- /dev/null +++ b/paimon-python/pypaimon/tests/paired_act_benchmark_test.py @@ -0,0 +1,355 @@ +# 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, + _git_head, + _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 + 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" + ).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) + + +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" diff --git a/paimon-python/setup.py b/paimon-python/setup.py index 9111730fb7ca..cccd2838e24b 100644 --- a/paimon-python/setup.py +++ b/paimon-python/setup.py @@ -243,6 +243,10 @@ def read_requirements(): 'torch': [ 'torch', ], + 'act': [ + 'lerobot==0.4.4', + 'Pillow', + ], 'daft': [ 'daft>=0.7.6; python_version>="3.10"', ],