From b7f2337613ea132bbcf64aba225db3bf84c1d1d1 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:01:18 +0200 Subject: [PATCH 01/65] Add Rerun reader --- pyproject.toml | 5 + src/refiner/__init__.py | 3 + src/refiner/pipeline/__init__.py | 3 + src/refiner/pipeline/pipeline.py | 58 +++ src/refiner/pipeline/sources/__init__.py | 2 + .../pipeline/sources/readers/__init__.py | 2 + src/refiner/pipeline/sources/readers/rerun.py | 393 ++++++++++++++++++ tests/readers/test_rerun_reader.py | 67 +++ uv.lock | 55 ++- 9 files changed, 587 insertions(+), 1 deletion(-) create mode 100644 src/refiner/pipeline/sources/readers/rerun.py create mode 100644 tests/readers/test_rerun_reader.py diff --git a/pyproject.toml b/pyproject.toml index 4de39281..12211ef1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,6 +64,10 @@ mcap = [ "mcap-ros2-support", "pillow", ] +rerun = [ + "pillow", + "rerun-sdk[datafusion]>=0.33,<0.34", +] s3 = [ "s3fs", ] @@ -87,6 +91,7 @@ all = [ "macrodata-refiner[hdf5]", "macrodata-refiner[hf]", "macrodata-refiner[mcap]", + "macrodata-refiner[rerun]", "macrodata-refiner[video]", "macrodata-refiner[zarr]", "macrodata-refiner[text]", diff --git a/src/refiner/__init__.py b/src/refiner/__init__.py index 15dba63f..b6ea9ede 100644 --- a/src/refiner/__init__.py +++ b/src/refiner/__init__.py @@ -31,6 +31,7 @@ "read_lerobot": "refiner.pipeline", "read_mcap": "refiner.pipeline", "read_parquet": "refiner.pipeline", + "read_rerun": "refiner.pipeline", "read_tfds": "refiner.pipeline", "read_tfrecords": "refiner.pipeline", "read_videos": "refiner.pipeline", @@ -66,6 +67,7 @@ "read_lerobot", "read_mcap", "read_parquet", + "read_rerun", "read_tfds", "read_tfrecords", "read_videos", @@ -140,6 +142,7 @@ def __dir__() -> list[str]: read_lerobot, read_mcap, read_parquet, + read_rerun, read_tfds, read_tfrecords, read_videos, diff --git a/src/refiner/pipeline/__init__.py b/src/refiner/pipeline/__init__.py index 856ec288..cf7389bd 100644 --- a/src/refiner/pipeline/__init__.py +++ b/src/refiner/pipeline/__init__.py @@ -18,6 +18,7 @@ "read_lerobot": "refiner.pipeline.pipeline", "read_mcap": "refiner.pipeline.pipeline", "read_parquet": "refiner.pipeline.pipeline", + "read_rerun": "refiner.pipeline.pipeline", "read_tfds": "refiner.pipeline.pipeline", "read_tfrecords": "refiner.pipeline.pipeline", "read_videos": "refiner.pipeline.pipeline", @@ -48,6 +49,7 @@ "read_lerobot", "read_mcap", "read_parquet", + "read_rerun", "read_tfds", "read_tfrecords", "read_videos", @@ -88,6 +90,7 @@ def __dir__() -> list[str]: read_lerobot, read_mcap, read_parquet, + read_rerun, read_tfds, read_tfrecords, read_videos, diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 4ee33797..daded149 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -41,6 +41,7 @@ JsonReader, McapReader, ParquetReader, + RerunReader, TfdsReader, TfrecordReader, ZarrReader, @@ -48,6 +49,7 @@ from refiner.pipeline.sources.readers.hdf5 import MissingPolicy from refiner.pipeline.sources.readers.lerobot import LeRobotEpisodeReader from refiner.pipeline.sources.readers.mcap import SyncMethod +from refiner.pipeline.sources.readers.rerun import RerunOutputMode from refiner.pipeline.sources.items import ItemsSource from refiner.pipeline.sources.task import TaskSource, TaskStep from refiner.pipeline.data import datatype @@ -1217,6 +1219,62 @@ def read_mcap( ) +def read_rerun( + inputs: DataFileSetLike, + *, + fs: AbstractFileSystem | None = None, + storage_options: Mapping[str, Any] | None = None, + recursive: bool = False, + target_shard_bytes: int = DEFAULT_TARGET_SHARD_BYTES, + num_shards: int | None = None, + file_path_column: str | None = "file_path", + output: RerunOutputMode = "recording", + contents: str | Sequence[str] | None = None, + timelines: Sequence[str] | None = None, + primary_timeline: str | None = None, + include_static: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + action_prefix: str = "/action", + state_prefix: str = "/observation/state", + camera_prefix: str = "/cam", + fps: float | None = None, + robot_type: str | None = None, +) -> RefinerPipeline: + """Create a pipeline with a Rerun RRD reader source. + + RRD files are planned as atomic input shards. With ``output="recording"``, + each emitted row preserves the selected Rerun data as Arrow-backed + ``Tabular`` tables grouped by timeline under the ``rerun`` field. With + ``output="robotics"``, the reader additionally derives common robotics + episode fields from configurable Rerun entity prefixes so the rows can be + passed through ``to_robot_rows(...)`` and robotics writers. + """ + return RefinerPipeline( + source=RerunReader( + inputs, + fs=fs, + storage_options=storage_options, + recursive=recursive, + target_shard_bytes=target_shard_bytes, + num_shards=num_shards, + file_path_column=file_path_column, + output=output, + contents=contents, + timelines=timelines, + primary_timeline=primary_timeline, + include_static=include_static, + include_recording=include_recording, + fill_latest_at=fill_latest_at, + action_prefix=action_prefix, + state_prefix=state_prefix, + camera_prefix=camera_prefix, + fps=fps, + robot_type=robot_type, + ) + ) + + def read_parquet( inputs: DataFileSetLike, *, diff --git a/src/refiner/pipeline/sources/__init__.py b/src/refiner/pipeline/sources/__init__.py index 3ef0c3e7..90857a84 100644 --- a/src/refiner/pipeline/sources/__init__.py +++ b/src/refiner/pipeline/sources/__init__.py @@ -9,6 +9,7 @@ LeRobotEpisodeReader, McapReader, ParquetReader, + RerunReader, TfdsReader, TfrecordReader, ZarrReader, @@ -25,6 +26,7 @@ "LeRobotEpisodeReader", "McapReader", "ParquetReader", + "RerunReader", "TfdsReader", "TfrecordReader", "ZarrReader", diff --git a/src/refiner/pipeline/sources/readers/__init__.py b/src/refiner/pipeline/sources/readers/__init__.py index 09a1184b..8a3eda84 100644 --- a/src/refiner/pipeline/sources/readers/__init__.py +++ b/src/refiner/pipeline/sources/readers/__init__.py @@ -7,6 +7,7 @@ from refiner.pipeline.sources.readers.lerobot import LeRobotEpisodeReader from refiner.pipeline.sources.readers.mcap import McapReader from refiner.pipeline.sources.readers.parquet import ParquetReader +from refiner.pipeline.sources.readers.rerun import RerunReader from refiner.pipeline.sources.readers.tfds import TfdsReader from refiner.pipeline.sources.readers.tfrecord import TfrecordReader from refiner.pipeline.sources.readers.zarr import ZarrReader @@ -23,6 +24,7 @@ "LeRobotRow", "McapReader", "ParquetReader", + "RerunReader", "TfdsReader", "TfrecordReader", "ZarrReader", diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py new file mode 100644 index 00000000..edb8da50 --- /dev/null +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -0,0 +1,393 @@ +from __future__ import annotations + +import os +import tempfile +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal, cast + +import numpy as np +import pyarrow as pa +from fsspec import AbstractFileSystem + +from refiner.io import DataFile +from refiner.io.fileset import DataFileSetLike +from refiner.pipeline.data.row import DictRow +from refiner.pipeline.data.shard import FilePartsDescriptor, Shard +from refiner.pipeline.data.tabular import Tabular +from refiner.pipeline.sources.base import SourceUnit +from refiner.pipeline.sources.readers.base import BaseReader +from refiner.pipeline.sources.readers.utils import DEFAULT_TARGET_SHARD_BYTES +from refiner.utils import check_required_dependencies +from refiner.video import VideoFrameSequence + +RerunOutputMode = Literal["recording", "robotics"] + +_INDEX_METADATA_KEY = b"rerun:kind" +_INDEX_METADATA_VALUE = b"index" +_RERUN_SEGMENT_ID = "rerun_segment_id" +_DEFAULT_ROBOTICS_CONTENTS = ("/action/**", "/observation/**", "/cam/**") + + +@dataclass(frozen=True, slots=True) +class RerunRecording: + """Columnar Rerun recording data loaded from one RRD segment.""" + + segment_id: str + source_path: str + tables: Mapping[str, Tabular] + static: Tabular | None = None + + +class RerunReader(BaseReader): + """Read Rerun RRD files as columnar recording rows or robotics episode rows.""" + + name = "read_rerun" + + def __init__( + self, + inputs: DataFileSetLike, + *, + fs: AbstractFileSystem | None = None, + storage_options: Mapping[str, Any] | None = None, + recursive: bool = False, + target_shard_bytes: int = DEFAULT_TARGET_SHARD_BYTES, + num_shards: int | None = None, + file_path_column: str | None = "file_path", + output: RerunOutputMode = "recording", + contents: str | Sequence[str] | None = None, + timelines: Sequence[str] | None = None, + primary_timeline: str | None = None, + include_static: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + action_prefix: str = "/action", + state_prefix: str = "/observation/state", + camera_prefix: str = "/cam", + fps: float | None = None, + robot_type: str | None = None, + ) -> None: + if output not in ("recording", "robotics"): + raise ValueError("output must be 'recording' or 'robotics'") + if fps is not None: + fps = float(fps) + if not np.isfinite(fps) or fps <= 0: + raise ValueError("fps must be > 0") + super().__init__( + inputs, + fs=fs, + storage_options=storage_options, + recursive=recursive, + extensions=(".rrd",), + target_shard_bytes=target_shard_bytes, + num_shards=num_shards, + file_path_column=file_path_column, + split_by_bytes=False, + ) + self.output = output + self.contents = _contents(contents) + self.timelines = tuple(timelines) if timelines is not None else None + self.primary_timeline = primary_timeline + self.include_static = include_static + self.include_recording = ( + output == "recording" if include_recording is None else include_recording + ) + self.fill_latest_at = fill_latest_at + self.action_prefix = _normalize_entity_prefix(action_prefix) + self.state_prefix = _normalize_entity_prefix(state_prefix) + self.camera_prefix = _normalize_entity_prefix(camera_prefix) + self.fps = fps + self.robot_type = robot_type + + def _declared_refiner_extras(self) -> tuple[str, ...]: + return ("rerun",) + + def describe(self) -> dict[str, Any]: + description = super().describe() + description.update( + { + "output": self.output, + "contents": self.contents, + "timelines": self.timelines, + "primary_timeline": self.primary_timeline, + "include_static": self.include_static, + "include_recording": self.include_recording, + "fill_latest_at": self.fill_latest_at, + } + ) + return description + + def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: + descriptor = shard.descriptor + assert isinstance(descriptor, FilePartsDescriptor) + for part in descriptor.parts: + source = self.fileset.resolve_file(part.source_index, part.path) + with _local_rrd(source) as local_path: + yield from self._read_file(source, local_path) + + def _read_file(self, source: DataFile, local_path: Path) -> Iterator[SourceUnit]: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk"), "datafusion"], + dist="rerun", + ) + import rerun as rr + + with rr.server.Server(datasets={"recording": [local_path]}) as server: + dataset = server.client().get_dataset("recording") + schema = dataset.schema() + timelines = self._timelines(schema) + static = ( + _collect_table(dataset.reader(index=None)) + if self.include_static + else None + ) + for segment_id in dataset.segment_ids(): + view = dataset.filter_segments([segment_id]) + if self.output == "robotics": + yield self._robotics_row( + view, + segment_id=segment_id, + source_path=source.abs_path(), + schema=schema, + timelines=timelines, + static=static, + ) + else: + tables = { + timeline: Tabular( + _collect_table( + self._view_for_contents(view).reader( + index=timeline, + fill_latest_at=self.fill_latest_at, + ) + ) + ) + for timeline in timelines + } + data: dict[str, Any] = { + "episode_id": segment_id, + "rerun": RerunRecording( + segment_id=segment_id, + source_path=source.abs_path(), + tables=tables, + static=Tabular(static) if static is not None else None, + ), + } + self._with_file_path(data, source) + yield DictRow(data) + + def _timelines(self, schema: Any) -> tuple[str, ...]: + if self.timelines is not None: + return self.timelines + return tuple(str(index.name) for index in schema.index_columns()) + + def _primary_timeline(self, timelines: Sequence[str]) -> str: + if self.primary_timeline is not None: + return self.primary_timeline + for timeline in timelines: + if timeline not in {"log_tick", "log_time", "real_time"}: + return timeline + if not timelines: + raise ValueError("Rerun recording has no timelines") + return timelines[0] + + def _view_for_contents(self, dataset_or_view: Any) -> Any: + if self.contents is None: + return dataset_or_view + return dataset_or_view.filter_contents(self.contents) + + def _robotics_row( + self, + view: Any, + *, + segment_id: str, + source_path: str, + schema: Any, + timelines: Sequence[str], + static: pa.Table | None, + ) -> DictRow: + timeline = self._primary_timeline(timelines) + contents = self.contents or _DEFAULT_ROBOTICS_CONTENTS + table = _collect_table( + view.filter_contents(contents).reader( + index=timeline, + fill_latest_at=self.fill_latest_at, + ) + ) + frames = _robotics_frame_table(table, timeline=timeline) + row: dict[str, Any] = { + "episode_id": segment_id, + } + if self.include_recording: + row["rerun"] = RerunRecording( + segment_id=segment_id, + source_path=source_path, + tables={timeline: Tabular(table)}, + static=Tabular(static) if static is not None else None, + ) + if self.fps is not None: + row["fps"] = self.fps + if self.robot_type is not None: + row["robot_type"] = self.robot_type + + scalar_columns = _component_columns( + schema, + component="Scalars:scalars", + table=table, + ) + action_columns = _prefixed_columns(scalar_columns, self.action_prefix) + state_columns = _prefixed_columns(scalar_columns, self.state_prefix) + if action_columns: + frames = frames.append_column( + "action", + _list_column(_singleton_scalar_matrix(table, action_columns)), + ) + if state_columns: + frames = frames.append_column( + "observation.state", + _list_column(_singleton_scalar_matrix(table, state_columns)), + ) + row["frames"] = Tabular(frames) + + for name, column in _camera_columns(schema, table, self.camera_prefix).items(): + values = table.column(column).combine_chunks() + row[name] = VideoFrameSequence( + lambda values=values: _iter_encoded_images(values), + fps=self.fps or 30.0, + frame_count=len(values), + ) + + if self.file_path_column is not None: + row[self.file_path_column] = source_path + return DictRow(row) + + +def _contents(contents: str | Sequence[str] | None) -> tuple[str, ...] | None: + if contents is None: + return None + if isinstance(contents, str): + return (contents,) + return tuple(contents) + + +def _normalize_entity_prefix(value: str) -> str: + value = value.strip() + if not value: + raise ValueError("Rerun entity prefixes must be non-empty") + return "/" + value.strip("/") + + +def _collect_table(df: Any) -> pa.Table: + table = df.to_arrow_table() + if _RERUN_SEGMENT_ID in table.column_names and table.num_rows > 0: + table = table.drop_null() + return table + + +class _local_rrd: + def __init__(self, source: DataFile) -> None: + self.source = source + self.tmpdir: tempfile.TemporaryDirectory[str] | None = None + self.path: Path | None = None + + def __enter__(self) -> Path: + if self.source.is_local: + return Path(self.source.abs_path()) + self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") + name = os.path.basename(self.source.path) or "recording.rrd" + self.path = Path(self.tmpdir.name) / name + self.source.copy(str(self.path)) + return self.path + + def __exit__(self, *args: object) -> None: + if self.tmpdir is not None: + self.tmpdir.cleanup() + + +def _component_columns( + schema: Any, + *, + component: str, + table: pa.Table, +) -> dict[str, str]: + out: dict[str, str] = {} + names = set(table.column_names) + for column in schema.component_columns(): + if str(column.component) != component: + continue + name = str(column.name) + if name in names: + out[str(column.entity_path)] = name + return out + + +def _prefixed_columns(columns: Mapping[str, str], prefix: str) -> list[tuple[str, str]]: + return sorted( + ((path, column) for path, column in columns.items() if path.startswith(prefix)), + key=lambda item: item[0], + ) + + +def _camera_columns(schema: Any, table: pa.Table, prefix: str) -> dict[str, str]: + names = set(table.column_names) + out: dict[str, str] = {} + for column in schema.component_columns(): + if str(column.component) != "EncodedImage:blob": + continue + entity_path = str(column.entity_path) + name = str(column.name) + if not entity_path.startswith(prefix) or name not in names: + continue + out[entity_path.strip("/").replace("/", ".")] = name + return out + + +def _robotics_frame_table(table: pa.Table, *, timeline: str) -> pa.Table: + columns: dict[str, pa.ChunkedArray] = {} + if timeline in table.column_names: + columns["frame_index"] = table.column(timeline) + return pa.table(columns) + + +def _singleton_scalar_matrix( + table: pa.Table, + columns: Sequence[tuple[str, str]], +) -> np.ndarray: + values = [] + for _, column in columns: + values.append(_singleton_list_array(table.column(column).combine_chunks())) + if not values: + return np.empty((table.num_rows, 0), dtype=np.float64) + return np.stack(values, axis=1) + + +def _list_column(values: np.ndarray) -> pa.Array: + return pa.array(values.tolist()) + + +def _singleton_list_array(array: pa.Array) -> np.ndarray: + out = np.full(len(array), np.nan, dtype=np.float64) + for index, scalar in enumerate(array): + value = scalar.as_py() + if value: + out[index] = float(value[0]) + return out + + +def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: + from io import BytesIO + + from PIL import Image + + for scalar in values: + value = scalar.as_py() + if not value: + continue + data = cast(bytes | bytearray | list[int], value[0]) + with Image.open(BytesIO(bytes(data))) as image: + yield np.asarray(image.convert("RGB"), dtype=np.uint8) + + +__all__ = ["RerunReader", "RerunRecording", "RerunOutputMode"] diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py new file mode 100644 index 00000000..fc14f928 --- /dev/null +++ b/tests/readers/test_rerun_reader.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Any, cast + +import numpy as np +import pytest + +import refiner as mdr +from refiner.pipeline import Row + +pytest.importorskip("rerun") + + +def _tiny_rrd(path: Path) -> None: + import rerun as rr + + rr.init("refiner_rerun_test", recording_id="episode-a") + rr.save(path) + frames = np.arange(3) + rr.send_columns( + "/action/x", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0, 3.0])), + ) + rr.send_columns( + "/observation/state/y", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([4.0, 5.0, 6.0])), + ) + + +def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + unit = next(mdr.read_rerun(str(rrd), timelines=("frame",)).source.read()) + assert isinstance(unit, Row) + row = cast(Any, unit) + recording = row["rerun"] + + assert row["episode_id"] == "episode-a" + assert list(recording.tables) == ["frame"] + assert recording.tables["frame"].num_rows == 3 + assert row["file_path"] == str(rrd) + + +def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0) + .to_robot_rows( + episode_id_key="episode_id", + nested_frames_key="frames", + fps=30.0, + ) + .take(1)[0], + ) + + assert "rerun" not in row + assert row.episode_id == "episode-a" + assert row.num_frames == 3 + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] diff --git a/uv.lock b/uv.lock index 9ade185d..06335d65 100644 --- a/uv.lock +++ b/uv.lock @@ -867,6 +867,23 @@ nvtx = [ { name = "nvidia-nvtx", marker = "sys_platform == 'linux'" }, ] +[[package]] +name = "datafusion" +version = "52.3.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pyarrow" }, + { name = "typing-extensions", marker = "python_full_version < '3.13'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/db/d4/a5ad7b665a80008901892fde61dc667318db0652a955d706ddca3a224b5a/datafusion-52.3.0.tar.gz", hash = "sha256:2e8b02ad142b1a0d673f035d96a0944a640ac78275003d7e453cee4afe4a20a4", size = 205026, upload-time = "2026-03-16T10:54:07.739Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/55/63/1bb0737988cefa77274b459d64fa4b57ba4cf755639a39733e9581b5d599/datafusion-52.3.0-cp310-abi3-macosx_10_12_x86_64.whl", hash = "sha256:a73f02406b2985b9145dd97f8221a929c9ef3289a8ba64c6b52043e240938528", size = 31503230, upload-time = "2026-03-16T10:53:50.312Z" }, + { url = "https://files.pythonhosted.org/packages/d6/e3/ea3b79239953c3044d19d8e9581015da025b6640796db03799e435b17910/datafusion-52.3.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:118a1f0add6a3f91fcbc90c71819fe08750e2981637d5e7b346e099e94a20b8b", size = 28159497, upload-time = "2026-03-16T10:53:54.032Z" }, + { url = "https://files.pythonhosted.org/packages/24/c8/7d325feb4b7509ae03857fd7e164e95ec72e8c9f3dfd3178ec7f80d53977/datafusion-52.3.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:253ce7aee5fe84bd6ee290c20608114114bdb5115852617f97d3855d36ad9341", size = 30769154, upload-time = "2026-03-16T10:53:57.835Z" }, + { url = "https://files.pythonhosted.org/packages/37/ee/478689c69b3cb1ccabb2d52feac0c181f6cdf20b51a81df35344b1dab9a6/datafusion-52.3.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:2af3469d2f06959bec88579ab107a72f965de18b32e607069bbdd0b859ed8dbb", size = 33060335, upload-time = "2026-03-16T10:54:01.715Z" }, + { url = "https://files.pythonhosted.org/packages/1c/48/01906ab5c1a70373c6874ac5192d03646fa7b94d9ff06e3f676cb6b0f43f/datafusion-52.3.0-cp310-abi3-win_amd64.whl", hash = "sha256:9fb35738cf4dbff672dbcfffc7332813024cb0ad2ab8cda1fb90b9054277ab0c", size = 33765807, upload-time = "2026-03-16T10:54:05.728Z" }, +] + [[package]] name = "datasets" version = "4.8.5" @@ -2085,6 +2102,7 @@ all = [ { name = "mcap-ros2-support" }, { name = "numcodecs" }, { name = "pillow" }, + { name = "rerun-sdk", extra = ["datafusion"] }, { name = "s3fs" }, { name = "tensorflow" }, { name = "tensorflow-datasets" }, @@ -2120,6 +2138,10 @@ mcap = [ { name = "mcap-ros2-support" }, { name = "pillow" }, ] +rerun = [ + { name = "pillow" }, + { name = "rerun-sdk", extra = ["datafusion"] }, +] s3 = [ { name = "s3fs" }, ] @@ -2140,6 +2162,7 @@ testing = [ { name = "pillow" }, { name = "pytest" }, { name = "pytest-cov" }, + { name = "rerun-sdk", extra = ["datafusion"] }, { name = "s3fs" }, { name = "tensorflow" }, { name = "tensorflow-datasets" }, @@ -2193,6 +2216,7 @@ requires-dist = [ { name = "macrodata-refiner", extras = ["hf"], marker = "extra == 'datasets'" }, { name = "macrodata-refiner", extras = ["hf"], marker = "extra == 'hand-tracking'" }, { name = "macrodata-refiner", extras = ["mcap"], marker = "extra == 'all'" }, + { name = "macrodata-refiner", extras = ["rerun"], marker = "extra == 'all'" }, { name = "macrodata-refiner", extras = ["s3"], marker = "extra == 'all'" }, { name = "macrodata-refiner", extras = ["tensorflow"], marker = "extra == 'tfds'" }, { name = "macrodata-refiner", extras = ["text"], marker = "extra == 'all'" }, @@ -2209,18 +2233,20 @@ requires-dist = [ { name = "orjson" }, { name = "packaging" }, { name = "pillow", marker = "extra == 'mcap'" }, + { name = "pillow", marker = "extra == 'rerun'" }, { name = "pillow", marker = "extra == 'video'" }, { name = "pyarrow" }, { name = "pydantic", specifier = ">=2.0.0" }, { name = "pytest", marker = "extra == 'testing'", specifier = ">=8.0.0" }, { name = "pytest-cov", marker = "extra == 'testing'", specifier = ">=5.0.0" }, + { name = "rerun-sdk", extras = ["datafusion"], marker = "extra == 'rerun'", specifier = ">=0.33,<0.34" }, { name = "s3fs", marker = "extra == 's3'" }, { name = "tensorflow", marker = "extra == 'tensorflow'" }, { name = "tensorflow-datasets", marker = "extra == 'tfds'" }, { name = "warcio", marker = "extra == 'text'" }, { name = "zarr", marker = "extra == 'zarr'", specifier = ">=2.18,<3" }, ] -provides-extras = ["video", "hf", "datasets", "hand-tracking", "text", "hdf5", "zarr", "mcap", "s3", "gcs", "tensorflow", "tfds", "testing", "all"] +provides-extras = ["video", "hf", "datasets", "hand-tracking", "text", "hdf5", "zarr", "mcap", "rerun", "s3", "gcs", "tensorflow", "tfds", "testing", "all"] [package.metadata.requires-dev] dev = [ @@ -4150,6 +4176,33 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3b/5d/63d4ae3b9daea098d5d6f5da83984853c1bbacd5dc826764b249fe119d24/requests_oauthlib-2.0.0-py2.py3-none-any.whl", hash = "sha256:7dd8a5c40426b779b0868c404bdef9768deccf22749cde15852df527e6269b36", size = 24179, upload-time = "2024-03-22T20:32:28.055Z" }, ] +[[package]] +name = "rerun-sdk" +version = "0.33.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "numpy", version = "2.2.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.4.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "pillow" }, + { name = "psutil" }, + { name = "pyarrow" }, + { name = "typing-extensions" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/31/17/5a521e86ac0064bd0f452e3e98e2422433511b54110423c0217d2cc1234f/rerun_sdk-0.33.0-cp310-abi3-macosx_11_0_arm64.whl", hash = "sha256:97f123e3ef6aa69b60194bc566e5435c7d4040757ed4f58297ea46c8ef320c5c", size = 125707606, upload-time = "2026-05-29T09:42:53.584Z" }, + { url = "https://files.pythonhosted.org/packages/34/2f/2ca2599aca03b69fbcac7c8391ef50376968edd7c58b96de53a4b7f20624/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:8f734cf59419dcfbc46915bea6cec030224f16e96c3a597f0ccf7cb7b058dd43", size = 135271020, upload-time = "2026-05-29T09:43:00.106Z" }, + { url = "https://files.pythonhosted.org/packages/2e/ba/d70997b43e6db4f58c4326c29c6a6a384ddc6c2fe125f231c885ad9b3b1f/rerun_sdk-0.33.0-cp310-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:53d95609f8b330026bcd041bf6d11b46ee1c18b6fbde155135f291fe86328eeb", size = 139552018, upload-time = "2026-05-29T09:43:06.275Z" }, + { url = "https://files.pythonhosted.org/packages/14/a5/0cac294d16aff6c9a2f183f838428a0380b4d2fd9e053bb37b3041999ad5/rerun_sdk-0.33.0-cp310-abi3-win_amd64.whl", hash = "sha256:b152992a72ec240062c8c285bd30ab681b464a25efbe1464c66fdac82320de1f", size = 120418186, upload-time = "2026-05-29T09:43:13.733Z" }, +] + +[package.optional-dependencies] +datafusion = [ + { name = "datafusion" }, + { name = "pandas", version = "2.3.3", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "pandas", version = "3.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, +] + [[package]] name = "rich" version = "14.3.3" From 917aee3d7b120d2e9490c834c2a4a5924de53be7 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:04:47 +0200 Subject: [PATCH 02/65] Add Rerun writer --- src/refiner/pipeline/pipeline.py | 26 +++- src/refiner/pipeline/sinks/__init__.py | 2 + src/refiner/pipeline/sinks/rerun.py | 192 +++++++++++++++++++++++++ tests/readers/test_rerun_reader.py | 22 +++ 4 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 src/refiner/pipeline/sinks/rerun.py diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index daded149..829c86ae 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -30,7 +30,7 @@ VectorizedSegmentStep, WithColumnsStep, ) -from refiner.pipeline.sinks import BaseSink, JsonlSink, ParquetSink, ZarrSink +from refiner.pipeline.sinks import BaseSink, JsonlSink, ParquetSink, RerunSink, ZarrSink from refiner.pipeline.sinks.assets import MissingAssetPolicy from refiner.pipeline.sources import ( BaseSource, @@ -619,6 +619,30 @@ def write_parquet( ) ) + def write_rerun( + self, + output: DataFolderLike, + *, + filename_template: str = "{shard_id}__w{worker_id}/{row_index}.rrd", + app_id: str = "refiner", + write_footer: bool = True, + ) -> "RefinerPipeline": + """Attach a distributed Rerun RRD writer sink. + + Rows must contain a ``RerunRecording`` value in the ``rerun`` field, + as emitted by ``read_rerun(output="recording")``. Each row is written + as one RRD file under ``output`` using a deterministic shard/worker + filename template. + """ + return self.with_sink( + RerunSink( + output=output, + filename_template=filename_template, + app_id=app_id, + write_footer=write_footer, + ) + ) + def write_zarr( self, output: DataFolderLike, diff --git a/src/refiner/pipeline/sinks/__init__.py b/src/refiner/pipeline/sinks/__init__.py index f0623f21..ccfb38f1 100644 --- a/src/refiner/pipeline/sinks/__init__.py +++ b/src/refiner/pipeline/sinks/__init__.py @@ -1,6 +1,7 @@ from refiner.pipeline.sinks.base import BaseSink, NullSink from refiner.pipeline.sinks.jsonl import JsonlSink from refiner.pipeline.sinks.parquet import ParquetSink +from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.reducer import FileCleanupReducerSink, LeRobotMetaReduceSink from refiner.pipeline.sinks.zarr import ZarrSink @@ -11,5 +12,6 @@ "JsonlSink", "LeRobotMetaReduceSink", "ParquetSink", + "RerunSink", "ZarrSink", ] diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py new file mode 100644 index 00000000..2520f0a7 --- /dev/null +++ b/src/refiner/pipeline/sinks/rerun.py @@ -0,0 +1,192 @@ +from __future__ import annotations + +import os +import tempfile +from pathlib import Path +from string import Formatter + +import pyarrow as pa + +from refiner.io.datafile import DataFile +from refiner.io.datafolder import DataFolder, DataFolderLike +from refiner.pipeline.data.block import Block +from refiner.pipeline.data.row import Row +from refiner.pipeline.sinks.base import BaseSink +from refiner.pipeline.sinks.reducer.file import FileCleanupReducerSink +from refiner.pipeline.sources.readers.rerun import RerunRecording +from refiner.utils import check_required_dependencies +from refiner.worker.context import get_active_worker_token + +_DEFAULT_FILENAME_TEMPLATE = "{shard_id}__w{worker_id}/{row_index}.rrd" + + +class RerunSink(BaseSink): + """Write Rerun recording rows as distributed shard-local RRD files.""" + + def __init__( + self, + output: DataFolderLike, + *, + filename_template: str = _DEFAULT_FILENAME_TEMPLATE, + app_id: str = "refiner", + write_footer: bool = True, + ) -> None: + _validate_filename_template(filename_template) + self.output = DataFolder.resolve(output) + self.filename_template = filename_template + self.app_id = app_id + self.write_footer = write_footer + self._row_indices: dict[str, int] = {} + + def _declared_refiner_extras(self) -> tuple[str, ...]: + return ("rerun",) + + def write_shard_block(self, shard_id: str, block: Block) -> int: + count = 0 + for row in block: + recording = _recording_from_row(row) + row_index = self._row_indices.get(shard_id, 0) + self._row_indices[shard_id] = row_index + 1 + relpath = _render_relpath( + self.filename_template, + shard_id=shard_id, + worker_id=get_active_worker_token(), + row_index=row_index, + segment_id=recording.segment_id, + ) + self._write_recording(recording, relpath) + count += 1 + return count + + def _write_recording(self, recording: RerunRecording, relpath: str) -> None: + check_required_dependencies( + "write_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + import rerun as rr + + def write_local(path: Path) -> None: + with rr.RecordingStream( + self.app_id, + recording_id=recording.segment_id, + ) as rec: + rec.save(path, write_footer=self.write_footer) + if recording.static is not None: + static = _sendable_table(recording.static.table) + if static.num_rows > 0 or static.num_columns > 0: + rec.send_dataframe(static) + for table in recording.tables.values(): + rec.send_dataframe(_sendable_table(table.table)) + + target = self.output.file(relpath) + if target.is_local: + local_path = Path(target.abs_path()) + local_path.parent.mkdir(parents=True, exist_ok=True) + write_local(local_path) + return + + with tempfile.TemporaryDirectory(prefix="refiner-rerun-write-") as tmpdir: + local_path = Path(tmpdir) / os.path.basename(relpath) + write_local(local_path) + DataFile.resolve(str(local_path)).copy(target) + + def on_shard_complete(self, shard_id: str) -> None: + self._row_indices.pop(shard_id, None) + + def describe(self) -> tuple[str, str, dict[str, object]]: + return ( + "write_rerun", + "writer", + { + "path": self.output.abs_path(), + "filename_template": self.filename_template, + "app_id": self.app_id, + "write_footer": self.write_footer, + }, + ) + + def build_reducer(self) -> BaseSink | None: + return FileCleanupReducerSink( + output=self.output, + filename_template=self.filename_template, + reducer_name="write_rerun_reduce", + ) + + +def _recording_from_row(row: Row) -> RerunRecording: + value = row.get("rerun") + if not isinstance(value, RerunRecording): + raise ValueError("write_rerun requires rows with a RerunRecording in 'rerun'") + return value + + +def _sendable_table(table: pa.Table) -> pa.Table: + if "rerun_segment_id" in table.column_names: + return table.drop(["rerun_segment_id"]) + return table + + +def _validate_filename_template(filename_template: str) -> None: + fields: set[str] = set() + for _literal_text, field_name, format_spec, conversion in Formatter().parse( + filename_template + ): + if field_name is None: + continue + if conversion is not None or format_spec: + raise ValueError("filename_template only supports plain named fields") + if field_name not in {"shard_id", "worker_id", "row_index", "segment_id"}: + raise ValueError( + "filename_template only supports shard_id, worker_id, " + "row_index, and segment_id" + ) + fields.add(field_name) + missing = {"shard_id", "worker_id"}.difference(fields) + if missing: + raise ValueError( + "filename_template requires fields: " + + ", ".join(f"{{{field}}}" for field in sorted(missing)) + ) + _normalize_relpath( + filename_template.format( + shard_id="shard", + worker_id="worker", + row_index=0, + segment_id="segment", + ), + "filename_template", + ) + + +def _render_relpath( + filename_template: str, + *, + shard_id: str, + worker_id: str, + row_index: int, + segment_id: str, +) -> str: + return _normalize_relpath( + filename_template.format( + shard_id=shard_id, + worker_id=worker_id, + row_index=row_index, + segment_id=segment_id, + ), + "rendered filename", + ) + + +def _normalize_relpath(path: str, label: str) -> str: + if path.startswith("/"): + raise ValueError(f"{label} must be relative") + parts = [part for part in path.split("/") if part] + if not parts: + raise ValueError(f"{label} must not be empty") + if any(part in {".", ".."} for part in parts): + raise ValueError(f"{label} must not contain '.' or '..' segments") + return "/".join(parts) + + +__all__ = ["RerunSink"] diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index fc14f928..f9cd7131 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -8,6 +8,7 @@ import refiner as mdr from refiner.pipeline import Row +from refiner.pipeline.sinks.rerun import RerunSink pytest.importorskip("rerun") @@ -65,3 +66,24 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: assert row.num_frames == 3 assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] + + +def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out" + _tiny_rrd(source) + + unit = next(mdr.read_rerun(str(source), timelines=("frame",)).source.read()) + assert isinstance(unit, Row) + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", [unit]) + sink.on_shard_complete("shard-a") + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + row = next(mdr.read_rerun(str(written[0]), timelines=("frame",)).source.read()) + + assert isinstance(row, Row) + recording = row["rerun"] + assert row["episode_id"] == "episode-a" + assert recording.tables["frame"].num_rows == 3 From 14c1d0109e61e7cf764dffd20db900a1de369101 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:14:33 +0200 Subject: [PATCH 03/65] Harden Rerun reader and writer --- src/refiner/pipeline/sinks/rerun.py | 178 +++++++++++++-- src/refiner/pipeline/sources/readers/rerun.py | 205 +++++++++++++----- tests/readers/test_rerun_reader.py | 79 +++++++ 3 files changed, 393 insertions(+), 69 deletions(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index 2520f0a7..b6b95035 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -4,6 +4,7 @@ import tempfile from pathlib import Path from string import Formatter +from typing import Any import pyarrow as pa @@ -64,20 +65,17 @@ def _write_recording(self, recording: RerunRecording, relpath: str) -> None: [("rerun", "rerun-sdk")], dist="rerun", ) - import rerun as rr def write_local(path: Path) -> None: - with rr.RecordingStream( - self.app_id, - recording_id=recording.segment_id, - ) as rec: - rec.save(path, write_footer=self.write_footer) - if recording.static is not None: - static = _sendable_table(recording.static.table) - if static.num_rows > 0 or static.num_columns > 0: - rec.send_dataframe(static) - for table in recording.tables.values(): - rec.send_dataframe(_sendable_table(table.table)) + if recording.use_source_chunks and recording.source_file is not None: + _write_source_chunks(recording, path, application_id=self.app_id) + return + _write_recording_tables( + recording, + path, + application_id=self.app_id, + write_footer=self.write_footer, + ) target = self.output.file(relpath) if target.is_local: @@ -121,10 +119,158 @@ def _recording_from_row(row: Row) -> RerunRecording: return value -def _sendable_table(table: pa.Table) -> pa.Table: - if "rerun_segment_id" in table.column_names: - return table.drop(["rerun_segment_id"]) - return table +def _write_source_chunks( + recording: RerunRecording, + path: Path, + *, + application_id: str, +) -> None: + import rerun as rr + + source = recording.source_file + if source is None: + raise ValueError("Rerun source chunk write requires source_file") + with _local_rrd(source) as local_path: + reader = rr.experimental.RrdReader(local_path) + store = _matching_store(reader, recording) + stream = reader.stream(store=store) + if recording.contents is not None: + stream = stream.filter(content=recording.contents) + if not recording.include_static: + stream = stream.drop(is_static=True) + stream = _filter_timelines( + stream, + reader=reader, + store=store, + recording=recording, + ) + stream.write_rrd( + path, + application_id=recording.application_id or application_id, + recording_id=recording.recording_id or recording.segment_id, + ) + + +def _filter_timelines( + stream: Any, + *, + reader: Any, + store: Any, + recording: RerunRecording, +) -> Any: + timelines = recording.timelines + if timelines is None: + return stream + if len(timelines) == 1: + dynamic = stream.filter(has_timeline=timelines[0]) + if not recording.include_static: + return dynamic + static = reader.stream(store=store).filter(is_static=True) + if recording.contents is not None: + static = static.filter(content=recording.contents) + import rerun as rr + + return rr.experimental.LazyChunkStream.merge(static, dynamic) + + selected = set(timelines) + + def keep_selected(chunk: Any) -> tuple[Any, ...]: + if chunk.is_static: + return (chunk,) if recording.include_static else () + return (chunk,) if selected.intersection(chunk.timeline_names) else () + + return stream.flat_map(keep_selected) + + +def _matching_store(reader: Any, recording: RerunRecording) -> Any: + stores = list(reader.recordings()) + if not stores: + return None + for store in stores: + if ( + recording.recording_id is not None + and store.recording_id == recording.recording_id + and ( + recording.application_id is None + or store.application_id == recording.application_id + ) + ): + return store + for store in stores: + if store.recording_id == recording.segment_id: + return store + return stores[0] + + +def _write_recording_tables( + recording: RerunRecording, + path: Path, + *, + application_id: str, + write_footer: bool, +) -> None: + import rerun as rr + + with rr.RecordingStream( + recording.application_id or application_id, + recording_id=recording.recording_id or recording.segment_id, + ) as rec: + rec.save(path, write_footer=write_footer) + if recording.static is not None: + static = _sendable_static_table(recording.static.table) + if static.num_columns > 0: + rec.send_dataframe(static) + for table in recording.tables.values(): + dynamic = _sendable_dynamic_table(table.table) + if dynamic.num_columns > 0: + rec.send_dataframe(dynamic) + + +def _sendable_static_table(table: pa.Table) -> pa.Table: + keep = [ + field.name + for field in table.schema + if _is_data_column(field) and _is_static_column(field) + ] + return table.select(keep) if keep else pa.table({}) + + +def _sendable_dynamic_table(table: pa.Table) -> pa.Table: + keep = [ + field.name + for field in table.schema + if field.name != "rerun_segment_id" and not _is_static_column(field) + ] + return table.select(keep) if keep else pa.table({}) + + +def _is_data_column(field: pa.Field) -> bool: + metadata = field.metadata or {} + return metadata.get(b"rerun:kind") == b"data" or b"rerun:entity_path" in metadata + + +def _is_static_column(field: pa.Field) -> bool: + return (field.metadata or {}).get(b"rerun:is_static") == b"true" + + +class _local_rrd: + def __init__(self, source: DataFile) -> None: + self.source = source + self.tmpdir: tempfile.TemporaryDirectory[str] | None = None + self.path: Path | None = None + + def __enter__(self) -> Path: + if self.source.is_local: + return Path(self.source.abs_path()) + self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-source-") + name = os.path.basename(self.source.path) or "recording.rrd" + self.path = Path(self.tmpdir.name) / name + self.source.copy(str(self.path)) + return self.path + + def __exit__(self, *args: object) -> None: + if self.tmpdir is not None: + self.tmpdir.cleanup() def _validate_filename_template(filename_template: str) -> None: diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index edb8da50..858ba754 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -2,6 +2,7 @@ import os import tempfile +from contextlib import ExitStack from collections.abc import Iterator, Mapping, Sequence from dataclasses import dataclass from pathlib import Path @@ -9,6 +10,7 @@ import numpy as np import pyarrow as pa +import pyarrow.compute as pc from fsspec import AbstractFileSystem from refiner.io import DataFile @@ -38,6 +40,13 @@ class RerunRecording: source_path: str tables: Mapping[str, Tabular] static: Tabular | None = None + source_file: DataFile | None = None + application_id: str | None = None + recording_id: str | None = None + contents: tuple[str, ...] | None = None + timelines: tuple[str, ...] | None = None + include_static: bool = True + use_source_chunks: bool = True class RerunReader(BaseReader): @@ -121,12 +130,17 @@ def describe(self) -> dict[str, Any]: def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: descriptor = shard.descriptor assert isinstance(descriptor, FilePartsDescriptor) - for part in descriptor.parts: - source = self.fileset.resolve_file(part.source_index, part.path) - with _local_rrd(source) as local_path: - yield from self._read_file(source, local_path) - - def _read_file(self, source: DataFile, local_path: Path) -> Iterator[SourceUnit]: + with ExitStack() as stack: + local_files = [] + for part in descriptor.parts: + source = self.fileset.resolve_file(part.source_index, part.path) + local_files.append((source, stack.enter_context(_local_rrd(source)))) + yield from self._read_files(local_files) + + def _read_files( + self, + local_files: Sequence[tuple[DataFile, Path]], + ) -> Iterator[SourceUnit]: check_required_dependencies( "read_rerun", [("rerun", "rerun-sdk"), "datafusion"], @@ -134,49 +148,80 @@ def _read_file(self, source: DataFile, local_path: Path) -> Iterator[SourceUnit] ) import rerun as rr - with rr.server.Server(datasets={"recording": [local_path]}) as server: - dataset = server.client().get_dataset("recording") - schema = dataset.schema() - timelines = self._timelines(schema) + datasets = { + f"recording_{index}": (str(local_path),) + for index, (_source, local_path) in enumerate(local_files) + } + with rr.server.Server(datasets=cast(Any, datasets)) as server: + client = server.client() + for dataset_name, (source, local_path) in zip( + datasets, local_files, strict=True + ): + dataset = client.get_dataset(dataset_name) + yield from self._read_dataset(source, local_path, dataset) + + def _read_dataset( + self, + source: DataFile, + local_path: Path, + dataset: Any, + ) -> Iterator[SourceUnit]: + store_entries = _recording_entries(local_path) + entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} + schema = dataset.schema() + timelines = self._timelines(schema) + for segment_id in dataset.segment_ids(): + store = entries_by_recording_id.get(segment_id) + application_id = store.application_id if store is not None else None + recording_id = store.recording_id if store is not None else segment_id + view = dataset.filter_segments([segment_id]) + content_view = self._view_for_contents(view) static = ( - _collect_table(dataset.reader(index=None)) + _collect_table(content_view.reader(index=None)) if self.include_static else None ) - for segment_id in dataset.segment_ids(): - view = dataset.filter_segments([segment_id]) - if self.output == "robotics": - yield self._robotics_row( - view, - segment_id=segment_id, - source_path=source.abs_path(), - schema=schema, - timelines=timelines, - static=static, - ) - else: - tables = { - timeline: Tabular( - _collect_table( - self._view_for_contents(view).reader( - index=timeline, - fill_latest_at=self.fill_latest_at, - ) + if self.output == "robotics": + yield self._robotics_row( + view, + segment_id=segment_id, + source_path=source.abs_path(), + source_file=source, + application_id=application_id, + recording_id=recording_id, + schema=schema, + timelines=timelines, + static=static, + ) + else: + tables = { + timeline: Tabular( + _collect_table( + content_view.reader( + index=timeline, + fill_latest_at=self.fill_latest_at, ) ) - for timeline in timelines - } - data: dict[str, Any] = { - "episode_id": segment_id, - "rerun": RerunRecording( - segment_id=segment_id, - source_path=source.abs_path(), - tables=tables, - static=Tabular(static) if static is not None else None, - ), - } - self._with_file_path(data, source) - yield DictRow(data) + ) + for timeline in timelines + } + data: dict[str, Any] = { + "episode_id": segment_id, + "rerun": RerunRecording( + segment_id=segment_id, + source_path=source.abs_path(), + tables=tables, + static=Tabular(static) if static is not None else None, + source_file=source, + application_id=application_id, + recording_id=recording_id, + contents=self.contents, + timelines=self.timelines, + include_static=self.include_static, + ), + } + self._with_file_path(data, source) + yield DictRow(data) def _timelines(self, schema: Any) -> tuple[str, ...]: if self.timelines is not None: @@ -204,6 +249,9 @@ def _robotics_row( *, segment_id: str, source_path: str, + source_file: DataFile, + application_id: str | None, + recording_id: str, schema: Any, timelines: Sequence[str], static: pa.Table | None, @@ -226,6 +274,12 @@ def _robotics_row( source_path=source_path, tables={timeline: Tabular(table)}, static=Tabular(static) if static is not None else None, + source_file=source_file, + application_id=application_id, + recording_id=recording_id, + contents=tuple(contents), + timelines=(timeline,), + include_static=self.include_static, ) if self.fps is not None: row["fps"] = self.fps @@ -282,10 +336,19 @@ def _normalize_entity_prefix(value: str) -> str: def _collect_table(df: Any) -> pa.Table: table = df.to_arrow_table() if _RERUN_SEGMENT_ID in table.column_names and table.num_rows > 0: - table = table.drop_null() + table = table.filter(_is_valid(table.column(_RERUN_SEGMENT_ID))) return table +def _recording_entries(local_path: Path) -> list[Any]: + import rerun as rr + + try: + return list(rr.experimental.RrdReader(local_path).recordings()) + except Exception: + return [] + + class _local_rrd: def __init__(self, source: DataFile) -> None: self.source = source @@ -369,10 +432,18 @@ def _list_column(values: np.ndarray) -> pa.Array: def _singleton_list_array(array: pa.Array) -> np.ndarray: out = np.full(len(array), np.nan, dtype=np.float64) - for index, scalar in enumerate(array): - value = scalar.as_py() - if value: - out[index] = float(value[0]) + if len(array) == 0: + return out + if not pa.types.is_list(array.type) and not pa.types.is_large_list(array.type): + raise TypeError(f"Expected a Rerun list component column, got {array.type}") + offsets = np.asarray(array.offsets) + starts = offsets[:-1] + ends = offsets[1:] + valid = np.asarray(_is_valid(array), dtype=bool) & (ends > starts) + if not valid.any(): + return out + values = np.asarray(array.values) + out[valid] = values[starts[valid]] return out @@ -381,13 +452,41 @@ def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: from PIL import Image - for scalar in values: - value = scalar.as_py() - if not value: + for index in range(len(values)): + data = _encoded_image_bytes(values, index) + if data is None: continue - data = cast(bytes | bytearray | list[int], value[0]) - with Image.open(BytesIO(bytes(data))) as image: + with Image.open(BytesIO(data)) as image: yield np.asarray(image.convert("RGB"), dtype=np.uint8) +def _encoded_image_bytes(values: pa.Array, index: int) -> bytes | None: + if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" + ) + if not values[index].is_valid: + return None + outer_offsets = np.asarray(values.offsets) + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + return None + inner = values.values + if not pa.types.is_list(inner.type) and not pa.types.is_large_list(inner.type): + value = values[index].as_py() + if not value: + return None + return bytes(cast(bytes | bytearray | list[int], value[0])) + inner_offsets = np.asarray(inner.offsets) + byte_start = int(inner_offsets[outer_start]) + byte_end = int(inner_offsets[outer_start + 1]) + payload = inner.values.slice(byte_start, byte_end - byte_start) + return np.asarray(payload).tobytes() + + +def _is_valid(values: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: + return pc.call_function("is_valid", [values]) + + __all__ = ["RerunReader", "RerunRecording", "RerunOutputMode"] diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index f9cd7131..438bf1f7 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -9,6 +9,7 @@ import refiner as mdr from refiner.pipeline import Row from refiner.pipeline.sinks.rerun import RerunSink +from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table pytest.importorskip("rerun") @@ -31,6 +32,28 @@ def _tiny_rrd(path: Path) -> None: ) +def _sparse_rrd(path: Path) -> None: + import rerun as rr + + rr.init("refiner_rerun_sparse_test", recording_id="episode-sparse") + rr.save(path) + rr.send_columns( + "/action/x", + indexes=[], + columns=rr.SeriesLines.columns(names=["x"]), + ) + rr.send_columns( + "/action/x", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0, 1, 2]))], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0, 3.0])), + ) + rr.send_columns( + "/action/y", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([1]))], + columns=rr.Scalars.columns(scalars=np.asarray([9.0])), + ) + + def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -46,6 +69,17 @@ def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: assert row["file_path"] == str(rrd) +def test_read_rerun_recording_preserves_sparse_rows(tmp_path: Path) -> None: + rrd = tmp_path / "sparse.rrd" + _sparse_rrd(rrd) + + row = cast(Any, next(mdr.read_rerun(str(rrd), timelines=("frame",)).source.read())) + table = row["rerun"].tables["frame"].table + + assert table.num_rows == 3 + assert table.column("frame").to_pylist() == [0, 1, 2] + + def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -68,6 +102,34 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] +def test_read_rerun_robotics_mode_writes_lerobot(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + out = tmp_path / "lerobot" + _tiny_rrd(rrd) + + ( + mdr.read_rerun(str(rrd), output="robotics", fps=30.0, robot_type="testbot") + .to_robot_rows( + episode_id_key="episode_id", + nested_frames_key="frames", + fps=30.0, + robot_type="testbot", + ) + .write_lerobot(str(out), max_video_prepare_in_flight=1) + .launch_local( + name="rerun-to-lerobot-test", + num_workers=1, + rundir=str(tmp_path / "run"), + ) + ) + + row = cast(Any, mdr.read_lerobot(str(out)).take(1)[0]) + assert row.episode_id == "episode-a" + assert row.num_frames == 3 + assert row.actions.to_pylist() == [[1.0], [2.0], [3.0]] + assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] + + def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: source = tmp_path / "tiny.rrd" output = tmp_path / "out" @@ -87,3 +149,20 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: recording = row["rerun"] assert row["episode_id"] == "episode-a" assert recording.tables["frame"].num_rows == 3 + + +def test_write_rerun_table_fallback_separates_static_columns(tmp_path: Path) -> None: + source = tmp_path / "sparse.rrd" + _sparse_rrd(source) + + row = cast( + Any, next(mdr.read_rerun(str(source), timelines=("frame",)).source.read()) + ) + recording = row["rerun"] + + static = _sendable_static_table(recording.static.table) + dynamic = _sendable_dynamic_table(recording.tables["frame"].table) + + assert "/action/x:SeriesLines:names" in static.column_names + assert "/action/x:SeriesLines:names" not in dynamic.column_names + assert "frame" in dynamic.column_names From 276bd9f5f4f3e4c233c9869cb3b62bf4cac34c62 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:22:12 +0200 Subject: [PATCH 04/65] Support explicit Rerun robotics selections --- src/refiner/pipeline/pipeline.py | 10 +- src/refiner/pipeline/sources/readers/rerun.py | 169 +++++++++++++++++- tests/readers/test_rerun_reader.py | 83 +++++++++ 3 files changed, 252 insertions(+), 10 deletions(-) diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 829c86ae..9ea9fbb8 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -1262,6 +1262,9 @@ def read_rerun( action_prefix: str = "/action", state_prefix: str = "/observation/state", camera_prefix: str = "/cam", + actions: PathSelection | None = None, + states: PathSelection | None = None, + videos: PathSelection | None = None, fps: float | None = None, robot_type: str | None = None, ) -> RefinerPipeline: @@ -1272,7 +1275,9 @@ def read_rerun( ``Tabular`` tables grouped by timeline under the ``rerun`` field. With ``output="robotics"``, the reader additionally derives common robotics episode fields from configurable Rerun entity prefixes so the rows can be - passed through ``to_robot_rows(...)`` and robotics writers. + passed through ``to_robot_rows(...)`` and robotics writers. Pass ``actions``, + ``states``, or ``videos`` to pin exact entity paths and output order instead + of using prefix-derived defaults. """ return RefinerPipeline( source=RerunReader( @@ -1293,6 +1298,9 @@ def read_rerun( action_prefix=action_prefix, state_prefix=state_prefix, camera_prefix=camera_prefix, + actions=actions, + states=states, + videos=videos, fps=fps, robot_type=robot_type, ) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 858ba754..3677cac4 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -20,7 +20,11 @@ from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.sources.base import SourceUnit from refiner.pipeline.sources.readers.base import BaseReader -from refiner.pipeline.sources.readers.utils import DEFAULT_TARGET_SHARD_BYTES +from refiner.pipeline.sources.readers.utils import ( + DEFAULT_TARGET_SHARD_BYTES, + PathSelection, + path_selection_map, +) from refiner.utils import check_required_dependencies from refiner.video import VideoFrameSequence @@ -29,7 +33,9 @@ _INDEX_METADATA_KEY = b"rerun:kind" _INDEX_METADATA_VALUE = b"index" _RERUN_SEGMENT_ID = "rerun_segment_id" -_DEFAULT_ROBOTICS_CONTENTS = ("/action/**", "/observation/**", "/cam/**") +_ROBOTICS_ROW_COLUMNS = frozenset( + {"episode_id", "rerun", "frames", "fps", "robot_type"} +) @dataclass(frozen=True, slots=True) @@ -74,6 +80,9 @@ def __init__( action_prefix: str = "/action", state_prefix: str = "/observation/state", camera_prefix: str = "/cam", + actions: PathSelection | None = None, + states: PathSelection | None = None, + videos: PathSelection | None = None, fps: float | None = None, robot_type: str | None = None, ) -> None: @@ -106,6 +115,35 @@ def __init__( self.action_prefix = _normalize_entity_prefix(action_prefix) self.state_prefix = _normalize_entity_prefix(state_prefix) self.camera_prefix = _normalize_entity_prefix(camera_prefix) + self.actions_explicit = actions is not None + self.states_explicit = states is not None + self.videos_explicit = videos is not None + self.actions = _selection_map( + actions, + format_name="Rerun actions", + derive_names_from_paths=False, + ) + self.states = _selection_map( + states, + format_name="Rerun states", + derive_names_from_paths=False, + ) + self.videos = _selection_map(videos, format_name="Rerun videos") + if output == "robotics": + if file_path_column in _ROBOTICS_ROW_COLUMNS: + raise ValueError( + f"file_path_column cannot use reserved Rerun robotics row " + f"column {file_path_column!r}" + ) + reserved_video_names = set(_ROBOTICS_ROW_COLUMNS) + if file_path_column is not None: + reserved_video_names.add(file_path_column) + video_collisions = set(self.videos).intersection(reserved_video_names) + if video_collisions: + raise ValueError( + "Rerun video output names cannot use reserved robotics row " + "columns: " + ", ".join(sorted(video_collisions)) + ) self.fps = fps self.robot_type = robot_type @@ -123,6 +161,12 @@ def describe(self) -> dict[str, Any]: "include_static": self.include_static, "include_recording": self.include_recording, "fill_latest_at": self.fill_latest_at, + "action_prefix": self.action_prefix, + "state_prefix": self.state_prefix, + "camera_prefix": self.camera_prefix, + "actions": dict(self.actions) if self.actions_explicit else None, + "states": dict(self.states) if self.states_explicit else None, + "videos": dict(self.videos) if self.videos_explicit else None, } ) return description @@ -243,6 +287,24 @@ def _view_for_contents(self, dataset_or_view: Any) -> Any: return dataset_or_view return dataset_or_view.filter_contents(self.contents) + def _robotics_contents(self) -> tuple[str, ...]: + if self.contents is not None: + return self.contents + contents: list[str] = [] + if self.actions_explicit: + contents.extend(self.actions.values()) + else: + contents.append(_prefix_contents(self.action_prefix)) + if self.states_explicit: + contents.extend(self.states.values()) + else: + contents.append(_prefix_contents(self.state_prefix)) + if self.videos_explicit: + contents.extend(self.videos.values()) + else: + contents.append(_prefix_contents(self.camera_prefix)) + return tuple(dict.fromkeys(contents)) + def _robotics_row( self, view: Any, @@ -257,7 +319,7 @@ def _robotics_row( static: pa.Table | None, ) -> DictRow: timeline = self._primary_timeline(timelines) - contents = self.contents or _DEFAULT_ROBOTICS_CONTENTS + contents = self._robotics_contents() table = _collect_table( view.filter_contents(contents).reader( index=timeline, @@ -291,8 +353,24 @@ def _robotics_row( component="Scalars:scalars", table=table, ) - action_columns = _prefixed_columns(scalar_columns, self.action_prefix) - state_columns = _prefixed_columns(scalar_columns, self.state_prefix) + action_columns = ( + _selected_columns( + scalar_columns, + self.actions, + format_name="Rerun action", + ) + if self.actions_explicit + else _prefixed_columns(scalar_columns, self.action_prefix) + ) + state_columns = ( + _selected_columns( + scalar_columns, + self.states, + format_name="Rerun state", + ) + if self.states_explicit + else _prefixed_columns(scalar_columns, self.state_prefix) + ) if action_columns: frames = frames.append_column( "action", @@ -305,7 +383,12 @@ def _robotics_row( ) row["frames"] = Tabular(frames) - for name, column in _camera_columns(schema, table, self.camera_prefix).items(): + camera_columns = ( + _selected_camera_columns(schema, table, self.videos) + if self.videos_explicit + else _camera_columns(schema, table, self.camera_prefix) + ) + for name, column in camera_columns.items(): values = table.column(column).combine_chunks() row[name] = VideoFrameSequence( lambda values=values: _iter_encoded_images(values), @@ -330,7 +413,28 @@ def _normalize_entity_prefix(value: str) -> str: value = value.strip() if not value: raise ValueError("Rerun entity prefixes must be non-empty") - return "/" + value.strip("/") + stripped = value.strip("/") + return "/" if not stripped else "/" + stripped + + +def _prefix_contents(prefix: str) -> str: + return "/**" if prefix == "/" else f"{prefix}/**" + + +def _selection_map( + value: PathSelection | None, + *, + format_name: str, + derive_names_from_paths: bool = True, +) -> dict[str, str]: + return { + name: _normalize_entity_prefix(path) + for name, path in path_selection_map( + value, + format_name=format_name, + derive_names_from_paths=derive_names_from_paths, + ).items() + } def _collect_table(df: Any) -> pa.Table: @@ -388,11 +492,36 @@ def _component_columns( def _prefixed_columns(columns: Mapping[str, str], prefix: str) -> list[tuple[str, str]]: return sorted( - ((path, column) for path, column in columns.items() if path.startswith(prefix)), + ( + (path, column) + for path, column in columns.items() + if _matches_entity_prefix(path, prefix) + ), key=lambda item: item[0], ) +def _selected_columns( + columns: Mapping[str, str], + selected: Mapping[str, str], + *, + format_name: str, +) -> list[tuple[str, str]]: + out: list[tuple[str, str]] = [] + for _name, path in selected.items(): + column = columns.get(path) + if column is None: + raise KeyError(f"{format_name} entity path not found: {path}") + out.append((path, column)) + return out + + +def _matches_entity_prefix(entity_path: str, prefix: str) -> bool: + if prefix == "/": + return entity_path.startswith("/") + return entity_path == prefix or entity_path.startswith(f"{prefix}/") + + def _camera_columns(schema: Any, table: pa.Table, prefix: str) -> dict[str, str]: names = set(table.column_names) out: dict[str, str] = {} @@ -401,12 +530,34 @@ def _camera_columns(schema: Any, table: pa.Table, prefix: str) -> dict[str, str] continue entity_path = str(column.entity_path) name = str(column.name) - if not entity_path.startswith(prefix) or name not in names: + if not _matches_entity_prefix(entity_path, prefix) or name not in names: continue out[entity_path.strip("/").replace("/", ".")] = name return out +def _selected_camera_columns( + schema: Any, + table: pa.Table, + selected: Mapping[str, str], +) -> dict[str, str]: + names = set(table.column_names) + by_entity_path: dict[str, str] = {} + for column in schema.component_columns(): + if str(column.component) != "EncodedImage:blob": + continue + name = str(column.name) + if name in names: + by_entity_path[str(column.entity_path)] = name + out: dict[str, str] = {} + for name, path in selected.items(): + column = by_entity_path.get(path) + if column is None: + raise KeyError(f"Rerun video entity path not found: {path}") + out[name] = column + return out + + def _robotics_frame_table(table: pa.Table, *, timeline: str) -> pa.Table: columns: dict[str, pa.ChunkedArray] = {} if timeline in table.column_names: diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 438bf1f7..a91153bf 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -1,5 +1,6 @@ from __future__ import annotations +from io import BytesIO from pathlib import Path from typing import Any, cast @@ -25,6 +26,11 @@ def _tiny_rrd(path: Path) -> None: indexes=[rr.TimeColumn("frame", sequence=frames)], columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0, 3.0])), ) + rr.send_columns( + "/action_extra/y", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([10.0, 20.0, 30.0])), + ) rr.send_columns( "/observation/state/y", indexes=[rr.TimeColumn("frame", sequence=frames)], @@ -54,6 +60,49 @@ def _sparse_rrd(path: Path) -> None: ) +def _custom_robotics_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_custom_robotics_test", recording_id="episode-custom") + rr.save(path) + frames = np.arange(2) + rr.send_columns( + "/robot/actions/z", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([30.0, 40.0])), + ) + rr.send_columns( + "/robot/actions/a", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([10.0, 20.0])), + ) + rr.send_columns( + "/robot/state/b", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0])), + ) + rr.send_columns( + "/robot/state/a", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=np.asarray([3.0, 4.0])), + ) + blobs: list[bytes] = [] + for color in ((1, 2, 3), (4, 5, 6)): + image = Image.new("RGB", (1, 1), color=color) + out = BytesIO() + image.save(out, format="PNG") + blobs.append(out.getvalue()) + rr.send_columns( + "/robot/cameras/top", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.EncodedImage.columns( + blob=blobs, + media_type=["image/png", "image/png"], + ), + ) + + def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -102,6 +151,40 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] +def test_read_rerun_robotics_mode_respects_explicit_selections( + tmp_path: Path, +) -> None: + rrd = tmp_path / "custom.rrd" + _custom_robotics_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + actions=("/robot/actions/z", "/robot/actions/a"), + states={ + "first": "/robot/state/a", + "second": "/robot/state/b", + }, + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1)[0], + ) + + assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row["frames"].column("observation.state").to_pylist() == [ + [3.0, 1.0], + [4.0, 2.0], + ] + video = row["observation.images.top"] + assert video.frame_count == 2 + assert [frame[0, 0].tolist() for frame in video.iter_frame_arrays()] == [ + [1, 2, 3], + [4, 5, 6], + ] + + def test_read_rerun_robotics_mode_writes_lerobot(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" out = tmp_path / "lerobot" From 0c580c941b10c8744b3c4705f4267edee2d348ca Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:24:52 +0200 Subject: [PATCH 05/65] Document Rerun reader and writer --- docs/nav.md | 2 + docs/reading-data/index.md | 1 + docs/reading-data/rerun.md | 115 ++++++++++++++++++++++++ docs/reference/optional-dependencies.md | 2 + docs/writing-data/index.md | 2 +- docs/writing-data/rerun.md | 58 ++++++++++++ 6 files changed, 179 insertions(+), 1 deletion(-) create mode 100644 docs/reading-data/rerun.md create mode 100644 docs/writing-data/rerun.md diff --git a/docs/nav.md b/docs/nav.md index 4b2ba88f..053aba63 100644 --- a/docs/nav.md +++ b/docs/nav.md @@ -20,6 +20,7 @@ - [HDF5](reading-data/hdf5.md) - [Zarr](reading-data/zarr.md) - [MCAP](reading-data/mcap.md) +- [Rerun](reading-data/rerun.md) - [Tabular files](reading-data/tabular-files.md) - [Files and videos](reading-data/files-and-videos.md) - [Hugging Face](reading-data/hugging-face.md) @@ -65,6 +66,7 @@ - [Writer model](writing-data/writer-model.md) - [LeRobot](writing-data/lerobot.md) - [Zarr](writing-data/zarr.md) +- [Rerun](writing-data/rerun.md) - [Parquet and JSONL](writing-data/parquet-and-jsonl.md) - [Media assets and reducers](writing-data/media-assets-and-reducers.md) diff --git a/docs/reading-data/index.md b/docs/reading-data/index.md index 62639d59..3a1f92ca 100644 --- a/docs/reading-data/index.md +++ b/docs/reading-data/index.md @@ -22,6 +22,7 @@ pipeline = mdr.read_lerobot("hf://datasets/lerobot/aloha_sim_transfer_cube_human | One HDF5 file per episode, or grouped HDF5 demos | `read_hdf5` | [HDF5](hdf5.md) | | Zarr replay buffer with episode boundaries | `read_zarr` | [Zarr](zarr.md) | | MCAP robotics or autonomy logs | `read_mcap` | [MCAP](mcap.md) | +| Rerun RRD recordings | `read_rerun` | [Rerun](rerun.md) | | Parquet, JSON, JSONL, CSV tables | `read_parquet`, `read_json`, `read_jsonl`, `read_csv` | [Tabular Files](tabular-files.md) | | Raw files or media files | `read_files`, `read_videos` | [Files and Videos](files-and-videos.md) | | Hugging Face datasets table | `read_hf_dataset` | [Hugging Face](hugging-face.md) | diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md new file mode 100644 index 00000000..d2440ab9 --- /dev/null +++ b/docs/reading-data/rerun.md @@ -0,0 +1,115 @@ +--- +title: "Rerun reader" +description: "Read Rerun RRD recordings as columnar or robotics episode rows" +--- + +# Rerun reader + +Use `read_rerun` for `.rrd` files written by Rerun. + +```python +import refiner as mdr + +pipeline = mdr.read_rerun( + "s3://bucket/run/**/*.rrd", + output="robotics", + fps=30, +) +``` + +Install `macrodata-refiner[rerun]` to use this reader. Add storage extras such +as `s3` when reading remote paths. + +Directory inputs are filtered to paths ending in `.rrd`. RRD files are planned +as atomic files, so workers parallelize across recordings instead of splitting +one recording by byte range. + +## Recording rows + +With `output="recording"`, `read_rerun` emits one row per Rerun recording +segment: + +```python +rows = mdr.read_rerun( + "/data/run/*.rrd", + output="recording", + contents=("/action/**", "/observation/**"), + timelines=("frame",), +) +``` + +Each row includes: + +| Column | Meaning | +| --- | --- | +| `episode_id` | Rerun recording id or segment id. | +| `rerun` | `RerunRecording` value with Arrow-backed `Tabular` tables by timeline. | +| `file_path` | Source RRD path, unless `file_path_column=None`. | + +`contents` is passed to Rerun's content filter. `timelines` limits the timeline +tables returned. If `timelines` is omitted, the reader materializes all timeline +indexes reported by the Rerun schema. + +## Robotics rows + +With `output="robotics"`, the reader creates rows that can be passed to +`to_robot_rows(...)` and robotics writers: + +```python +robot_rows = ( + mdr.read_rerun( + "/data/episodes/*.rrd", + output="robotics", + fps=30, + robot_type="unknown", + ) + .to_robot_rows( + episode_id_key="episode_id", + nested_frames_key="frames", + fps_key="fps", + robot_type_key="robot_type", + video_keys={ + "observation.images.top": "cam.top", + "observation.images.left_wrist": "cam.left_wrist", + }, + ) +) +``` + +The default robotics mapping reads scalar components under `/action/**` into +the frame `action` vector, scalar components under `/observation/state/**` into +`observation.state`, and encoded images under `/cam/**` into top-level video +sources such as `cam.top`. + +Use explicit selections when vector order or camera names matter: + +```python +mdr.read_rerun( + "episode.rrd", + output="robotics", + actions=("/robot/actions/gripper", "/robot/actions/arm"), + states=("/robot/state/qpos", "/robot/state/gripper"), + videos={"observation.images.top": "/robot/cameras/top"}, + fps=30, +) +``` + +`actions` and `states` define vector order. `videos` maps output video keys to +Rerun encoded-image entity paths. If `contents` is omitted, explicit selections +also define the minimal Rerun content filter for those categories. + +## Decoding + +Scalar action and state columns are read from Arrow list arrays and converted +to frame vectors. Encoded images remain lazy `VideoFrameSequence` values; JPEG +or PNG bytes are decoded frame-by-frame only when a downstream video writer or +consumer iterates the sequence. + +## Sharding + +Rerun SDK queries require a complete local RRD file. For that reason, +`read_rerun` sets file-atomic sharding, like other container readers such as +HDF5 and MCAP. `target_shard_bytes` groups whole RRD files into shard buckets, +and `num_shards` can request a target number of file buckets when there are +enough files. + diff --git a/docs/reference/optional-dependencies.md b/docs/reference/optional-dependencies.md index 399d632b..cb77e1a3 100644 --- a/docs/reference/optional-dependencies.md +++ b/docs/reference/optional-dependencies.md @@ -16,6 +16,7 @@ Install extras based on the data and operations you use. | `hdf5` | HDF5 reader support. | | `zarr` | Zarr reader and writer support. | | `mcap` | MCAP robotics log reader support, including ROS2, protobuf, and H.264 video decoding. | +| `rerun` | Rerun RRD reader and writer support. | | `video` | Video decode/write support. | | `text` | Common Crawl text readers. | | `s3` | S3 filesystem support. | @@ -30,5 +31,6 @@ pip install macrodata-refiner[hf,video] pip install macrodata-refiner[datasets] pip install macrodata-refiner[hdf5,zarr] pip install macrodata-refiner[mcap] +pip install macrodata-refiner[rerun] pip install macrodata-refiner[hand_tracking] ``` diff --git a/docs/writing-data/index.md b/docs/writing-data/index.md index c2b0b8a0..ced3fc0c 100644 --- a/docs/writing-data/index.md +++ b/docs/writing-data/index.md @@ -12,6 +12,7 @@ reader, transforms, and writer stages. | --- | --- | | [LeRobot](lerobot.md) | Training-ready robotics datasets. | | [Zarr](zarr.md) | Array stores and replay buffers. | +| [Rerun](rerun.md) | Distributed `.rrd` recording outputs. | | [Parquet and JSONL](parquet-and-jsonl.md) | Tabular outputs and logs. | | [Media Assets and Reducers](media-assets-and-reducers.md) | Asset uploads, video handling, and reducer stages. | @@ -26,4 +27,3 @@ pipeline = ( ``` The writer does work when the pipeline is launched. - diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md new file mode 100644 index 00000000..7f212e52 --- /dev/null +++ b/docs/writing-data/rerun.md @@ -0,0 +1,58 @@ +--- +title: "Rerun writer" +description: "Write Rerun recording rows as distributed RRD files" +--- + +# Rerun writer + +Use `write_rerun` to write rows containing a `RerunRecording` value, as emitted +by `read_rerun(output="recording")`. + +```python +pipeline = ( + mdr.read_rerun( + "s3://bucket/input/*.rrd", + output="recording", + contents=("/action/**", "/observation/**"), + ) + .write_rerun("s3://bucket/output/rrd") +) +``` + +Install `macrodata-refiner[rerun]` to use this writer. Add storage extras such +as `s3` when writing remote paths. + +## Output layout + +`write_rerun` writes one RRD file per input recording row. The default file name +template is: + +```text +{shard_id}__w{worker_id}/{row_index}.rrd +``` + +The template must include `{shard_id}` and `{worker_id}` so retry cleanup can +distinguish finalized worker outputs from abandoned attempt outputs. You can +also use `{row_index}` and `{segment_id}`. + +## Writer strategy + +When the input row came from `read_rerun`, the writer uses Rerun's raw +`LazyChunkStream` path and writes the selected source chunks directly. This +preserves Rerun chunk metadata and avoids re-emitting large Arrow tables through +Python. + +If a `RerunRecording` has no source file, the writer falls back to table +emission with `send_dataframe`. Static Rerun component columns are sent as +static data, and dynamic timeline tables are sent separately. + +## Reducer + +The writer is distributed. Each worker writes deterministic shard-local files, +then a reducer stage removes files from non-finalized worker attempts. There is +no global merge step because the output is a directory of independent RRD +recordings. + +Use `write_lerobot` instead when the goal is a single training-ready robotics +dataset with merged LeRobot metadata. + From d6faf32eade412f283586fcbe40f24539b1b4e3e Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:29:24 +0200 Subject: [PATCH 06/65] Clean up Rerun docs whitespace --- docs/reading-data/rerun.md | 1 - docs/writing-data/rerun.md | 1 - 2 files changed, 2 deletions(-) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index d2440ab9..f783c973 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -112,4 +112,3 @@ Rerun SDK queries require a complete local RRD file. For that reason, HDF5 and MCAP. `target_shard_bytes` groups whole RRD files into shard buckets, and `num_shards` can request a target number of file buckets when there are enough files. - diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index 7f212e52..d7d534fb 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -55,4 +55,3 @@ recordings. Use `write_lerobot` instead when the goal is a single training-ready robotics dataset with merged LeRobot metadata. - From 087e2d3c051ec8512fbc2969e45e3d66e7894978 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:30:34 +0200 Subject: [PATCH 07/65] Avoid unused Rerun static reads --- src/refiner/pipeline/sources/readers/rerun.py | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 3677cac4..9fb5fd50 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -220,11 +220,6 @@ def _read_dataset( recording_id = store.recording_id if store is not None else segment_id view = dataset.filter_segments([segment_id]) content_view = self._view_for_contents(view) - static = ( - _collect_table(content_view.reader(index=None)) - if self.include_static - else None - ) if self.output == "robotics": yield self._robotics_row( view, @@ -235,9 +230,13 @@ def _read_dataset( recording_id=recording_id, schema=schema, timelines=timelines, - static=static, ) else: + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) tables = { timeline: Tabular( _collect_table( @@ -316,7 +315,6 @@ def _robotics_row( recording_id: str, schema: Any, timelines: Sequence[str], - static: pa.Table | None, ) -> DictRow: timeline = self._primary_timeline(timelines) contents = self._robotics_contents() @@ -331,6 +329,11 @@ def _robotics_row( "episode_id": segment_id, } if self.include_recording: + static = ( + _collect_table(view.filter_contents(contents).reader(index=None)) + if self.include_static + else None + ) row["rerun"] = RerunRecording( segment_id=segment_id, source_path=source_path, From aa805588ede244509b56014634fea464bae5bce2 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 15:38:46 +0200 Subject: [PATCH 08/65] Polish Rerun reader and writer --- docs/writing-data/rerun.md | 4 ++- src/refiner/pipeline/sinks/rerun.py | 6 ++++- src/refiner/pipeline/sources/readers/rerun.py | 19 ++++++++++--- tests/readers/test_rerun_reader.py | 27 +++++++++++++++++++ 4 files changed, 50 insertions(+), 6 deletions(-) diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index d7d534fb..70fa576e 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -44,7 +44,9 @@ Python. If a `RerunRecording` has no source file, the writer falls back to table emission with `send_dataframe`. Static Rerun component columns are sent as -static data, and dynamic timeline tables are sent separately. +static data, and dynamic timeline tables are sent separately. The same fallback +is used when `write_footer=False`, because Rerun's raw chunk writer always +writes footer metadata. ## Reducer diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index b6b95035..7d440bf3 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -67,7 +67,11 @@ def _write_recording(self, recording: RerunRecording, relpath: str) -> None: ) def write_local(path: Path) -> None: - if recording.use_source_chunks and recording.source_file is not None: + if ( + self.write_footer + and recording.use_source_chunks + and recording.source_file is not None + ): _write_source_chunks(recording, path, application_id=self.app_id) return _write_recording_tables( diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 9fb5fd50..c8b8891e 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -30,8 +30,6 @@ RerunOutputMode = Literal["recording", "robotics"] -_INDEX_METADATA_KEY = b"rerun:kind" -_INDEX_METADATA_VALUE = b"index" _RERUN_SEGMENT_ID = "rerun_segment_id" _ROBOTICS_ROW_COLUMNS = frozenset( {"episode_id", "rerun", "frames", "fps", "robot_type"} @@ -219,7 +217,6 @@ def _read_dataset( application_id = store.application_id if store is not None else None recording_id = store.recording_id if store is not None else segment_id view = dataset.filter_segments([segment_id]) - content_view = self._view_for_contents(view) if self.output == "robotics": yield self._robotics_row( view, @@ -232,6 +229,7 @@ def _read_dataset( timelines=timelines, ) else: + content_view = self._view_for_contents(view) static = ( _collect_table(content_view.reader(index=None)) if self.include_static @@ -581,7 +579,20 @@ def _singleton_scalar_matrix( def _list_column(values: np.ndarray) -> pa.Array: - return pa.array(values.tolist()) + if values.ndim != 2: + raise ValueError("Rerun vector columns must be 2D") + width = int(values.shape[1]) + if width <= 0: + offsets = pa.array(np.zeros(values.shape[0] + 1, dtype=np.int32)) + return pa.ListArray.from_arrays(offsets, pa.array([], type=pa.float64())) + flat_values = pa.array( + np.ascontiguousarray(values).reshape(-1), + type=pa.float64(), + ) + offsets = pa.array( + np.arange(0, len(flat_values) + width, width, dtype=np.int32), + ) + return pa.ListArray.from_arrays(offsets, flat_values) def _singleton_list_array(array: pa.Array) -> np.ndarray: diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index a91153bf..cd2ad755 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -234,6 +234,33 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: assert recording.tables["frame"].num_rows == 3 +def test_write_rerun_without_footer_uses_table_fallback( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-no-footer" + _tiny_rrd(source) + + unit = next(mdr.read_rerun(str(source), timelines=("frame",)).source.read()) + assert isinstance(unit, Row) + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun._write_source_chunks", + lambda *args, **kwargs: pytest.fail("raw chunk writer cannot disable footers"), + ) + sink = RerunSink(str(output), write_footer=False) + sink.write_shard_block("shard-a", [unit]) + sink.on_shard_complete("shard-a") + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + row = next(mdr.read_rerun(str(written[0]), timelines=("frame",)).source.read()) + + assert isinstance(row, Row) + assert row["rerun"].tables["frame"].num_rows == 3 + + def test_write_rerun_table_fallback_separates_static_columns(tmp_path: Path) -> None: source = tmp_path / "sparse.rrd" _sparse_rrd(source) From 012167a81c07b06c9f422ea4237751751028fc3c Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 22:42:11 +0200 Subject: [PATCH 09/65] Add Rerun cloud benchmark harness --- benchmark/rerun/README.md | 55 +++ benchmark/rerun/run_cloud_benchmark.py | 544 +++++++++++++++++++++++++ 2 files changed, 599 insertions(+) create mode 100644 benchmark/rerun/README.md create mode 100644 benchmark/rerun/run_cloud_benchmark.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md new file mode 100644 index 00000000..5a866a8b --- /dev/null +++ b/benchmark/rerun/README.md @@ -0,0 +1,55 @@ +# Rerun Benchmarks + +This folder contains cloud benchmark harnesses for Rerun RRD workloads. + +- `run_cloud_benchmark.py`: submits Macrodata Cloud jobs for Rerun read, + robotics conversion, and Rerun write paths, then writes JSON artifacts with + job ids, stage timings, metrics, and output inspection. + +The default inputs are the ten base RRD files from: + +```text +s3://macrodata-rerun-format-tests/dominique-sample/ +``` + +The default cases are: + +- `recording-summary`: `read_rerun(output="recording")`, summarize timeline + and static tables, write JSONL. +- `robotics-summary`: `read_rerun(output="robotics")` for action/state paths, + summarize frame rows and vector widths, write JSONL. +- `rrd-copy`: `read_rerun(output="recording").write_rerun(...)` to exercise the + distributed RRD writer's raw chunk path. + +These cases intentionally cover both the high-fidelity recording path and the +robotics convenience path. Do not remove a case just to make a performance run +look better. + +## Prerequisites + +- The current branch must be pushed and available on a GitHub PR before cloud + launch. +- Macrodata CLI auth must be configured. +- Workspace secrets in the selected environment must include AWS credentials + for the source/output S3 bucket. The default environment is `researcher`. + +## Run + +```bash +REFINER_ATTACH=detach uv run python benchmark/rerun/run_cloud_benchmark.py +``` + +Useful options: + +- `--case robotics-summary --case rrd-copy` to run a subset. +- `--iterations 3` to repeat each case. +- `--input s3://bucket/path/file.rrd` to use custom inputs; repeat as needed. +- `--output-root s3://bucket/prefix` to choose where cloud outputs are written. +- `--num-workers 4` to vary cloud parallelism. +- `--aws-profile 210049840512_Researcher` to inspect S3 outputs locally with a + specific profile after cloud completion. + +Artifacts are written under `benchmark/rerun/artifacts/` by default: + +- one per-case result JSON +- one summary JSON for the benchmark session diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py new file mode 100644 index 00000000..53d6a5bc --- /dev/null +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -0,0 +1,544 @@ +from __future__ import annotations + +import argparse +import json +import os +import platform +import re +import subprocess +import sys +import time +from collections.abc import Sequence +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +from fsspec import url_to_fs + +import refiner as mdr +from refiner.pipeline.data.row import DictRow, Row +from refiner.platform.client import MacrodataClient + +DEFAULT_INPUTS = tuple( + f"s3://macrodata-rerun-format-tests/dominique-sample/episode-{index}__base.rrd" + for index in range(10) +) +DEFAULT_OUTPUT_ROOT = "s3://macrodata-rerun-format-tests/refiner-rerun-benchmark" +DEFAULT_ARTIFACTS_DIR = Path(__file__).resolve().parent / "artifacts" +DEFAULT_CASES = ("recording-summary", "robotics-summary", "rrd-copy") +AWS_SECRET_KEYS = ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_DEFAULT_REGION", +) +TERMINAL_STATUSES = {"completed", "failed", "cancelled", "canceled"} + + +@dataclass(slots=True) +class StageResult: + index: int + name: str + status: str + n_shards: int | None + shard_done: int | None + shard_total: int | None + requested_workers: int | None + cpu_cores: int | None + memory_mb: int | None + duration_s: float | None + metrics: dict[str, dict[str, float | int | str | None]] + + +@dataclass(slots=True) +class CaseResult: + case: str + iteration: int + job_id: str + status: str + started_at_utc: str + finished_at_utc: str + input_count: int + output_root: str + cloud_wall_time_s: float | None + queue_time_s: float | None + stage_results: list[StageResult] + output_file_count: int | None + output_size_bytes: int | None + output_inspection_error: str | None + python_version: str + platform: str + git_ref: str + package_versions: dict[str, str] + + +def summarize_recording(row: Row) -> DictRow: + recording = row["rerun"] + table_summaries = { + name: { + "rows": table.table.num_rows, + "columns": table.table.num_columns, + "bytes": int(table.table.nbytes), + } + for name, table in recording.tables.items() + } + static = recording.static.table if recording.static is not None else None + return DictRow( + { + "episode_id": row["episode_id"], + "table_count": len(table_summaries), + "tables": table_summaries, + "static_columns": static.num_columns if static is not None else 0, + "static_bytes": int(static.nbytes) if static is not None else 0, + "application_id": recording.application_id, + "recording_id": recording.recording_id, + }, + shard_id=row.shard_id, + ) + + +def summarize_robotics(row: Row) -> DictRow: + table = row["frames"].table + action = table.column("action") if "action" in table.column_names else None + state = ( + table.column("observation.state") + if "observation.state" in table.column_names + else None + ) + return DictRow( + { + "episode_id": row["episode_id"], + "num_frames": table.num_rows, + "frame_columns": table.column_names, + "action_type": ( + str(table.schema.field("action").type) if action is not None else None + ), + "state_type": ( + str(table.schema.field("observation.state").type) + if state is not None + else None + ), + "first_action_width": ( + len(action[0].as_py() or []) + if action is not None and table.num_rows + else 0 + ), + "first_state_width": ( + len(state[0].as_py() or []) + if state is not None and table.num_rows + else 0 + ), + }, + shard_id=row.shard_id, + ) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run Macrodata Cloud benchmarks for Rerun reader/writer paths." + ) + parser.add_argument( + "--input", + dest="inputs", + action="append", + help=( + "Input RRD file, directory, or glob. Repeat for multiple inputs. " + "Defaults to the ten Dominique sample base RRDs." + ), + ) + parser.add_argument( + "--case", + dest="cases", + action="append", + choices=DEFAULT_CASES, + help=( + "Benchmark case to run. Repeat for multiple cases. Defaults to all " + f"cases: {', '.join(DEFAULT_CASES)}." + ), + ) + parser.add_argument("--iterations", type=int, default=1) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument("--cpus-per-worker", type=int, default=4) + parser.add_argument("--mem-mb-per-worker", type=int, default=8192) + parser.add_argument("--timeline", default="frame") + parser.add_argument("--fps", type=float, default=30.0) + parser.add_argument("--secret-env", default="researcher") + parser.add_argument("--output-root", default=DEFAULT_OUTPUT_ROOT) + parser.add_argument("--artifacts-dir", type=Path, default=DEFAULT_ARTIFACTS_DIR) + parser.add_argument("--run-token") + parser.add_argument("--poll-interval-s", type=float, default=10.0) + parser.add_argument("--timeout-s", type=float, default=60.0 * 60.0) + parser.add_argument( + "--aws-profile", + help=( + "Optional AWS profile used by local output inspection after the " + "cloud job completes." + ), + ) + parser.add_argument( + "--skip-output-inspection", + action="store_true", + help="Do not inspect output object counts/sizes from the submitting machine.", + ) + return parser.parse_args() + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _package_version(name: str) -> str: + try: + from importlib.metadata import version + + return version(name) + except Exception: + return "unknown" + + +def _package_versions() -> dict[str, str]: + return { + "macrodata-refiner": _package_version("macrodata-refiner"), + "rerun-sdk": _package_version("rerun-sdk"), + "datafusion": _package_version("datafusion"), + "pyarrow": _package_version("pyarrow"), + "s3fs": _package_version("s3fs"), + } + + +def _git_ref() -> str: + try: + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parents[2], + text=True, + stderr=subprocess.DEVNULL, + ).strip() + except Exception: + return "unknown" + + +def _sanitize_segment(value: str) -> str: + sanitized = re.sub(r"[^a-zA-Z0-9._-]+", "-", value).strip("-") + if not sanitized: + raise ValueError("path segment cannot be empty") + return sanitized + + +def _output_for_case( + *, + output_root: str, + run_token: str, + case: str, + iteration: int, +) -> str: + return "/".join( + [ + output_root.rstrip("/"), + _sanitize_segment(run_token), + _sanitize_segment(case), + f"iteration-{iteration:02d}", + ] + ) + + +def _build_pipeline( + *, + case: str, + inputs: Sequence[str], + output: str, + timeline: str, + fps: float, +) -> mdr.RefinerPipeline: + if case == "recording-summary": + return ( + mdr.read_rerun(inputs, output="recording", timelines=(timeline,)) + .map(summarize_recording) + .write_jsonl(output) + ) + if case == "robotics-summary": + return ( + mdr.read_rerun( + inputs, + output="robotics", + contents=("/action/**", "/observation/state/**"), + timelines=(timeline,), + include_recording=False, + fps=fps, + ) + .map(summarize_robotics) + .write_jsonl(output) + ) + if case == "rrd-copy": + return mdr.read_rerun(inputs, output="recording").write_rerun(output) + raise ValueError(f"Unsupported benchmark case: {case}") + + +def _wait_for_job( + client: MacrodataClient, + *, + job_id: str, + poll_interval_s: float, + timeout_s: float, +) -> dict[str, Any]: + deadline = time.monotonic() + timeout_s + while True: + payload = client.cli_get_job(job_id=job_id) + status = str(payload.get("status") or "") + if status in TERMINAL_STATUSES: + return payload + if time.monotonic() > deadline: + raise TimeoutError(f"Timed out waiting for cloud job {job_id}") + time.sleep(max(1.0, poll_interval_s)) + + +def _duration_s(started_ms: Any, ended_ms: Any) -> float | None: + if not isinstance(started_ms, (int, float)) or not isinstance( + ended_ms, (int, float) + ): + return None + return max(0.0, (float(ended_ms) - float(started_ms)) / 1000.0) + + +def _metric_values( + client: MacrodataClient, + *, + job_id: str, + stage_index: int, + step_index: int, + labels: Sequence[str], +) -> dict[str, dict[str, float | int | str | None]]: + payload = client.cli_get_job_step_metrics( + job_id=job_id, + stage_index=stage_index, + step_index=step_index, + metric_labels=list(labels), + ) + steps = payload.get("steps") + if not isinstance(steps, list) or not steps: + return {} + metrics = steps[0].get("metrics") + if not isinstance(metrics, list): + return {} + out: dict[str, dict[str, float | int | str | None]] = {} + for metric in metrics: + if not isinstance(metric, dict): + continue + label = metric.get("label") + if not isinstance(label, str): + continue + out[label] = { + "total": metric.get("total"), + "rate_since_start": metric.get("rateSinceStart"), + "per_worker": metric.get("perWorker"), + "unit": metric.get("unit"), + } + return out + + +def _stage_results(client: MacrodataClient, job: dict[str, Any]) -> list[StageResult]: + job_id = str(job["id"]) + stages = job.get("stages") + if not isinstance(stages, list): + return [] + out: list[StageResult] = [] + for stage in stages: + if not isinstance(stage, dict): + continue + stage_index = int(stage.get("index", 0)) + metrics: dict[str, dict[str, float | int | str | None]] = {} + steps = stage.get("steps") + if isinstance(steps, list): + for step in steps: + if not isinstance(step, dict): + continue + step_index = step.get("index") + if not isinstance(step_index, int): + continue + metrics.update( + _metric_values( + client, + job_id=job_id, + stage_index=stage_index, + step_index=step_index, + labels=( + "rows_read", + "rows_processed", + "rows_written", + "files_written", + ), + ) + ) + runtime = stage.get("runtimeConfig") + runtime = runtime if isinstance(runtime, dict) else {} + out.append( + StageResult( + index=stage_index, + name=str(stage.get("name") or ""), + status=str(stage.get("status") or ""), + n_shards=_optional_int(stage.get("nShards")), + shard_done=_optional_int(stage.get("shardDone")), + shard_total=_optional_int(stage.get("shardTotal")), + requested_workers=_optional_int(runtime.get("requestedNumWorkers")), + cpu_cores=_optional_int(runtime.get("cpuCores")), + memory_mb=_optional_int(runtime.get("memoryMb")), + duration_s=_duration_s(stage.get("startedAt"), stage.get("endedAt")), + metrics=metrics, + ) + ) + return out + + +def _optional_int(value: Any) -> int | None: + return int(value) if isinstance(value, (int, float)) else None + + +def _inspect_output(path: str) -> tuple[int | None, int | None, str | None]: + try: + fs, fs_path = url_to_fs(path) + if not fs.exists(fs_path): + return 0, 0, None + total_size = 0 + total_files = 0 + for child in fs.find(fs_path): + info = fs.info(child) + if info.get("type") == "directory": + continue + total_files += 1 + total_size += int(info.get("size", 0)) + return total_files, total_size, None + except Exception as err: + return None, None, str(err) + + +def _run_case( + *, + args: argparse.Namespace, + client: MacrodataClient, + case: str, + iteration: int, + inputs: Sequence[str], + git_ref: str, + run_token: str, +) -> CaseResult: + output = _output_for_case( + output_root=args.output_root, + run_token=run_token, + case=case, + iteration=iteration, + ) + pipeline = _build_pipeline( + case=case, + inputs=inputs, + output=output, + timeline=args.timeline, + fps=args.fps, + ) + started_at = _utc_now() + os.environ.setdefault("REFINER_ATTACH", "detach") + launch = pipeline.launch_cloud( + name=f"rerun-benchmark-{case}-{iteration:02d}-{git_ref[:8]}", + num_workers=args.num_workers, + cpus_per_worker=args.cpus_per_worker, + mem_mb_per_worker=args.mem_mb_per_worker, + secrets=mdr.Secrets.env(name=args.secret_env, keys=AWS_SECRET_KEYS), + ) + job = _wait_for_job( + client, + job_id=launch.job_id, + poll_interval_s=args.poll_interval_s, + timeout_s=args.timeout_s, + ) + finished_at = _utc_now() + output_file_count: int | None = None + output_size_bytes: int | None = None + output_error: str | None = None + if not args.skip_output_inspection: + output_file_count, output_size_bytes, output_error = _inspect_output(output) + + return CaseResult( + case=case, + iteration=iteration, + job_id=launch.job_id, + status=str(job.get("status") or ""), + started_at_utc=started_at, + finished_at_utc=finished_at, + input_count=len(inputs), + output_root=output, + cloud_wall_time_s=_duration_s(job.get("startedAt"), job.get("endedAt")), + queue_time_s=_duration_s(job.get("createdAt"), job.get("startedAt")), + stage_results=_stage_results(client, job), + output_file_count=output_file_count, + output_size_bytes=output_size_bytes, + output_inspection_error=output_error, + python_version=sys.version.replace("\n", " "), + platform=platform.platform(), + git_ref=git_ref, + package_versions=_package_versions(), + ) + + +def _write_result(path: Path, payload: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + +def main() -> int: + args = _parse_args() + if args.iterations < 1: + raise ValueError("--iterations must be >= 1") + if args.num_workers < 1: + raise ValueError("--num-workers must be >= 1") + if args.aws_profile: + os.environ["AWS_PROFILE"] = args.aws_profile + inputs = tuple(args.inputs or DEFAULT_INPUTS) + cases = tuple(args.cases or DEFAULT_CASES) + git_ref = _git_ref() + run_token = args.run_token or ( + datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + f"-{git_ref[:8]}" + ) + client = MacrodataClient() + results: list[CaseResult] = [] + for iteration in range(args.iterations): + for case in cases: + print(f"Running {case} iteration {iteration}...", flush=True) + result = _run_case( + args=args, + client=client, + case=case, + iteration=iteration, + inputs=inputs, + git_ref=git_ref, + run_token=run_token, + ) + results.append(result) + result_path = ( + args.artifacts_dir / run_token / f"{case}-{iteration:02d}.json" + ) + _write_result(result_path, asdict(result)) + print( + f"Finished {case} iteration {iteration}: " + f"{result.status} job={result.job_id} " + f"cloud_wall_time_s={result.cloud_wall_time_s}", + flush=True, + ) + + summary = { + "run_token": run_token, + "git_ref": git_ref, + "inputs": list(inputs), + "cases": list(cases), + "iterations": args.iterations, + "results": [asdict(result) for result in results], + } + summary_path = args.artifacts_dir / run_token / "summary.json" + _write_result(summary_path, summary) + print(f"Summary written to {summary_path}") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 4e7b2f6689687df063a4f597a9b8d2d9e4171d58 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 22:50:08 +0200 Subject: [PATCH 10/65] Harden Rerun cloud benchmark harness --- benchmark/rerun/.gitignore | 2 + benchmark/rerun/README.md | 3 ++ benchmark/rerun/run_cloud_benchmark.py | 58 +++++++++++++++++++++----- 3 files changed, 53 insertions(+), 10 deletions(-) create mode 100644 benchmark/rerun/.gitignore diff --git a/benchmark/rerun/.gitignore b/benchmark/rerun/.gitignore new file mode 100644 index 00000000..70cfce69 --- /dev/null +++ b/benchmark/rerun/.gitignore @@ -0,0 +1,2 @@ +artifacts/ +__pycache__/ diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 5a866a8b..0c7a4e36 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -48,6 +48,9 @@ Useful options: - `--num-workers 4` to vary cloud parallelism. - `--aws-profile 210049840512_Researcher` to inspect S3 outputs locally with a specific profile after cloud completion. +- `--continue-on-failure` to keep launching later cases after one case fails. + By default the harness records the failed case and stops, so bad credentials or + setup failures do not create a misleading benchmark session. Artifacts are written under `benchmark/rerun/artifacts/` by default: diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index 53d6a5bc..0fcadf38 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -57,6 +57,7 @@ class CaseResult: iteration: int job_id: str status: str + job_error: str | None started_at_utc: str finished_at_utc: str input_count: int @@ -181,6 +182,11 @@ def _parse_args() -> argparse.Namespace: action="store_true", help="Do not inspect output object counts/sizes from the submitting machine.", ) + parser.add_argument( + "--continue-on-failure", + action="store_true", + help="Continue running later cases after a cloud job fails.", + ) return parser.parse_args() @@ -462,6 +468,7 @@ def _run_case( iteration=iteration, job_id=launch.job_id, status=str(job.get("status") or ""), + job_error=job.get("error") if isinstance(job.get("error"), str) else None, started_at_utc=started_at, finished_at_utc=finished_at, input_count=len(inputs), @@ -486,6 +493,28 @@ def _write_result(path: Path, payload: Any) -> None: ) +def _write_summary( + *, + args: argparse.Namespace, + run_token: str, + git_ref: str, + inputs: Sequence[str], + cases: Sequence[str], + results: Sequence[CaseResult], +) -> Path: + summary = { + "run_token": run_token, + "git_ref": git_ref, + "inputs": list(inputs), + "cases": list(cases), + "iterations": args.iterations, + "results": [asdict(result) for result in results], + } + summary_path = args.artifacts_dir / run_token / "summary.json" + _write_result(summary_path, summary) + return summary_path + + def main() -> int: args = _parse_args() if args.iterations < 1: @@ -525,17 +554,26 @@ def main() -> int: f"cloud_wall_time_s={result.cloud_wall_time_s}", flush=True, ) + if result.status != "completed" and not args.continue_on_failure: + summary_path = _write_summary( + args=args, + run_token=run_token, + git_ref=git_ref, + inputs=inputs, + cases=cases, + results=results, + ) + print(f"Summary written to {summary_path}") + return 1 - summary = { - "run_token": run_token, - "git_ref": git_ref, - "inputs": list(inputs), - "cases": list(cases), - "iterations": args.iterations, - "results": [asdict(result) for result in results], - } - summary_path = args.artifacts_dir / run_token / "summary.json" - _write_result(summary_path, summary) + summary_path = _write_summary( + args=args, + run_token=run_token, + git_ref=git_ref, + inputs=inputs, + cases=cases, + results=results, + ) print(f"Summary written to {summary_path}") return 0 From 2331a1f197b158ad71681eb5c354d5ebb682406a Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 22:50:21 +0200 Subject: [PATCH 11/65] Skip unused Rerun metadata reads --- src/refiner/pipeline/sources/readers/rerun.py | 6 ++++- tests/readers/test_rerun_reader.py | 23 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index c8b8891e..f331add4 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -208,7 +208,11 @@ def _read_dataset( local_path: Path, dataset: Any, ) -> Iterator[SourceUnit]: - store_entries = _recording_entries(local_path) + store_entries = ( + _recording_entries(local_path) + if self.output == "recording" or self.include_recording + else [] + ) entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} schema = dataset.schema() timelines = self._timelines(schema) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index cd2ad755..47180037 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -151,6 +151,29 @@ def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: assert row.states.to_pylist() == [[4.0], [5.0], [6.0]] +def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun._recording_entries", + lambda *args, **kwargs: pytest.fail( + "robotics rows without recording payload do not need store metadata" + ), + ) + + row = cast( + Any, + mdr.read_rerun(str(rrd), output="robotics", fps=30.0).take(1)[0], + ) + + assert "rerun" not in row + assert row["frames"].num_rows == 3 + + def test_read_rerun_robotics_mode_respects_explicit_selections( tmp_path: Path, ) -> None: From 98f5f7b7870c31ed5afbf5957f2a6f8daa6f802d Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 22:59:21 +0200 Subject: [PATCH 12/65] Avoid redundant Rerun schema work --- src/refiner/pipeline/sources/readers/rerun.py | 66 +++++++++---------- tests/readers/test_rerun_reader.py | 39 +++++++++++ 2 files changed, 70 insertions(+), 35 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index f331add4..c41c1570 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -31,6 +31,8 @@ RerunOutputMode = Literal["recording", "robotics"] _RERUN_SEGMENT_ID = "rerun_segment_id" +_RERUN_COMPONENT_METADATA = b"rerun:component" +_RERUN_ENTITY_PATH_METADATA = b"rerun:entity_path" _ROBOTICS_ROW_COLUMNS = frozenset( {"episode_id", "rerun", "frames", "fps", "robot_type"} ) @@ -214,8 +216,9 @@ def _read_dataset( else [] ) entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} - schema = dataset.schema() - timelines = self._timelines(schema) + timelines = self.timelines + if timelines is None: + timelines = self._timelines(dataset.schema()) for segment_id in dataset.segment_ids(): store = entries_by_recording_id.get(segment_id) application_id = store.application_id if store is not None else None @@ -229,7 +232,6 @@ def _read_dataset( source_file=source, application_id=application_id, recording_id=recording_id, - schema=schema, timelines=timelines, ) else: @@ -315,7 +317,6 @@ def _robotics_row( source_file: DataFile, application_id: str | None, recording_id: str, - schema: Any, timelines: Sequence[str], ) -> DictRow: timeline = self._primary_timeline(timelines) @@ -354,9 +355,8 @@ def _robotics_row( row["robot_type"] = self.robot_type scalar_columns = _component_columns( - schema, - component="Scalars:scalars", table=table, + component="Scalars:scalars", ) action_columns = ( _selected_columns( @@ -389,9 +389,9 @@ def _robotics_row( row["frames"] = Tabular(frames) camera_columns = ( - _selected_camera_columns(schema, table, self.videos) + _selected_camera_columns(table, self.videos) if self.videos_explicit - else _camera_columns(schema, table, self.camera_prefix) + else _camera_columns(table, self.camera_prefix) ) for name, column in camera_columns.items(): values = table.column(column).combine_chunks() @@ -445,7 +445,9 @@ def _selection_map( def _collect_table(df: Any) -> pa.Table: table = df.to_arrow_table() if _RERUN_SEGMENT_ID in table.column_names and table.num_rows > 0: - table = table.filter(_is_valid(table.column(_RERUN_SEGMENT_ID))) + segment_ids = table.column(_RERUN_SEGMENT_ID) + if segment_ids.null_count: + table = table.filter(_is_valid(segment_ids)) return table @@ -478,20 +480,24 @@ def __exit__(self, *args: object) -> None: self.tmpdir.cleanup() +def _metadata_text(metadata: Mapping[bytes, bytes], key: bytes) -> str | None: + value = metadata.get(key) + return value.decode("utf-8") if value is not None else None + + def _component_columns( - schema: Any, + table: pa.Table, *, component: str, - table: pa.Table, ) -> dict[str, str]: out: dict[str, str] = {} - names = set(table.column_names) - for column in schema.component_columns(): - if str(column.component) != component: + for field in table.schema: + metadata = field.metadata or {} + if _metadata_text(metadata, _RERUN_COMPONENT_METADATA) != component: continue - name = str(column.name) - if name in names: - out[str(column.entity_path)] = name + entity_path = _metadata_text(metadata, _RERUN_ENTITY_PATH_METADATA) + if entity_path is not None: + out[entity_path] = field.name return out @@ -527,33 +533,23 @@ def _matches_entity_prefix(entity_path: str, prefix: str) -> bool: return entity_path == prefix or entity_path.startswith(f"{prefix}/") -def _camera_columns(schema: Any, table: pa.Table, prefix: str) -> dict[str, str]: - names = set(table.column_names) +def _camera_columns(table: pa.Table, prefix: str) -> dict[str, str]: out: dict[str, str] = {} - for column in schema.component_columns(): - if str(column.component) != "EncodedImage:blob": + for entity_path, column in _component_columns( + table, + component="EncodedImage:blob", + ).items(): + if not _matches_entity_prefix(entity_path, prefix): continue - entity_path = str(column.entity_path) - name = str(column.name) - if not _matches_entity_prefix(entity_path, prefix) or name not in names: - continue - out[entity_path.strip("/").replace("/", ".")] = name + out[entity_path.strip("/").replace("/", ".")] = column return out def _selected_camera_columns( - schema: Any, table: pa.Table, selected: Mapping[str, str], ) -> dict[str, str]: - names = set(table.column_names) - by_entity_path: dict[str, str] = {} - for column in schema.component_columns(): - if str(column.component) != "EncodedImage:blob": - continue - name = str(column.name) - if name in names: - by_entity_path[str(column.entity_path)] = name + by_entity_path = _component_columns(table, component="EncodedImage:blob") out: dict[str, str] = {} for name, path in selected.items(): column = by_entity_path.get(path) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 47180037..d02d0a8a 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -208,6 +208,45 @@ def test_read_rerun_robotics_mode_respects_explicit_selections( ] +def test_read_rerun_robotics_mode_with_explicit_timeline_uses_table_metadata( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "custom.rrd" + _custom_robotics_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun.RerunReader._timelines", + lambda *args, **kwargs: pytest.fail( + "explicit timelines should not require schema timeline discovery" + ), + ) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + timelines=("frame",), + actions=("/robot/actions/z", "/robot/actions/a"), + states={ + "first": "/robot/state/a", + "second": "/robot/state/b", + }, + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1)[0], + ) + + assert row["frames"].column("action").to_pylist() == [[30.0, 10.0], [40.0, 20.0]] + assert row["frames"].column("observation.state").to_pylist() == [ + [3.0, 1.0], + [4.0, 2.0], + ] + video = row["observation.images.top"] + assert video.frame_count == 2 + + def test_read_rerun_robotics_mode_writes_lerobot(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" out = tmp_path / "lerobot" From 85c4ccd0ce544f0b39d0f3ed22b0317b480e8449 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 23:01:58 +0200 Subject: [PATCH 13/65] Cover Rerun robotics recording payloads --- tests/readers/test_rerun_reader.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index d02d0a8a..ad35b3c7 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -174,6 +174,30 @@ def test_read_rerun_robotics_mode_without_recording_skips_recording_entries( assert row["frames"].num_rows == 3 +def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + mdr.read_rerun( + str(rrd), + output="robotics", + include_recording=True, + timelines=("frame",), + fps=30.0, + ).take(1)[0], + ) + + recording = row["rerun"] + assert recording.recording_id == "episode-a" + assert list(recording.tables) == ["frame"] + assert recording.tables["frame"].num_rows == 3 + assert row["frames"].num_rows == 3 + + def test_read_rerun_robotics_mode_respects_explicit_selections( tmp_path: Path, ) -> None: From 0e6ec55b9dc4d72d6e14fdb7a7e2512926f7a0df Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 23:12:32 +0200 Subject: [PATCH 14/65] Add Rerun benchmark comparison helper --- benchmark/rerun/README.md | 15 ++ benchmark/rerun/compare_results.py | 233 +++++++++++++++++++++++++++++ 2 files changed, 248 insertions(+) create mode 100644 benchmark/rerun/compare_results.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 0c7a4e36..4ef36903 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -5,6 +5,8 @@ This folder contains cloud benchmark harnesses for Rerun RRD workloads. - `run_cloud_benchmark.py`: submits Macrodata Cloud jobs for Rerun read, robotics conversion, and Rerun write paths, then writes JSON artifacts with job ids, stage timings, metrics, and output inspection. +- `compare_results.py`: compares two benchmark `summary.json` artifacts and + prints case-level and stage-level timing deltas. The default inputs are the ten base RRD files from: @@ -56,3 +58,16 @@ Artifacts are written under `benchmark/rerun/artifacts/` by default: - one per-case result JSON - one summary JSON for the benchmark session + +## Compare + +After running a baseline and candidate benchmark, compare their summaries: + +```bash +uv run python benchmark/rerun/compare_results.py \ + benchmark/rerun/artifacts/baseline/summary.json \ + benchmark/rerun/artifacts/candidate/summary.json +``` + +Only completed jobs are used for timing deltas. Failed jobs still appear in the +run-count columns so setup problems are visible instead of silently averaged in. diff --git a/benchmark/rerun/compare_results.py b/benchmark/rerun/compare_results.py new file mode 100644 index 00000000..43019c86 --- /dev/null +++ b/benchmark/rerun/compare_results.py @@ -0,0 +1,233 @@ +from __future__ import annotations + +import argparse +import json +import statistics +from collections import defaultdict +from collections.abc import Iterable, Mapping, Sequence +from pathlib import Path +from typing import Any + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Compare two Rerun cloud benchmark summary artifacts." + ) + parser.add_argument("baseline", type=Path, help="Baseline summary.json") + parser.add_argument("candidate", type=Path, help="Candidate summary.json") + parser.add_argument( + "--json", + action="store_true", + help="Print machine-readable JSON instead of tables.", + ) + return parser.parse_args() + + +def _load_summary(path: Path) -> dict[str, Any]: + with path.open("r", encoding="utf-8") as handle: + payload = json.load(handle) + if not isinstance(payload, dict): + raise ValueError(f"{path} is not a benchmark summary object") + results = payload.get("results") + if not isinstance(results, list): + raise ValueError(f"{path} does not contain a results list") + return payload + + +def _completed(results: Iterable[Mapping[str, Any]]) -> list[Mapping[str, Any]]: + return [result for result in results if result.get("status") == "completed"] + + +def _mean(values: Iterable[float | int | None]) -> float | None: + numbers = [float(value) for value in values if isinstance(value, (int, float))] + if not numbers: + return None + return statistics.fmean(numbers) + + +def _group_results(summary: Mapping[str, Any]) -> dict[str, list[Mapping[str, Any]]]: + grouped: dict[str, list[Mapping[str, Any]]] = defaultdict(list) + for result in summary.get("results", []): + if not isinstance(result, dict): + continue + case = result.get("case") + if isinstance(case, str): + grouped[case].append(result) + return dict(grouped) + + +def _stage_key(stage: Mapping[str, Any]) -> str: + name = stage.get("name") + if isinstance(name, str) and name: + return name + return f"stage-{stage.get('index', '?')}" + + +def _stage_means(results: Sequence[Mapping[str, Any]]) -> dict[str, float | None]: + durations: dict[str, list[float | None]] = defaultdict(list) + for result in results: + stages = result.get("stage_results") + if not isinstance(stages, list): + continue + for stage in stages: + if not isinstance(stage, dict): + continue + durations[_stage_key(stage)].append(stage.get("duration_s")) + return {stage: _mean(values) for stage, values in durations.items()} + + +def _delta( + baseline: float | None, + candidate: float | None, +) -> tuple[float | None, float | None]: + if baseline is None or candidate is None: + return None, None + absolute = candidate - baseline + percent = (absolute / baseline * 100.0) if baseline else None + return absolute, percent + + +def _comparison( + baseline: Mapping[str, Any], + candidate: Mapping[str, Any], +) -> dict[str, Any]: + baseline_grouped = _group_results(baseline) + candidate_grouped = _group_results(candidate) + cases = sorted(set(baseline_grouped) | set(candidate_grouped)) + rows: list[dict[str, Any]] = [] + for case in cases: + baseline_all = baseline_grouped.get(case, []) + candidate_all = candidate_grouped.get(case, []) + baseline_completed = _completed(baseline_all) + candidate_completed = _completed(candidate_all) + baseline_wall = _mean( + result.get("cloud_wall_time_s") for result in baseline_completed + ) + candidate_wall = _mean( + result.get("cloud_wall_time_s") for result in candidate_completed + ) + wall_delta_s, wall_delta_pct = _delta(baseline_wall, candidate_wall) + baseline_stages = _stage_means(baseline_completed) + candidate_stages = _stage_means(candidate_completed) + stage_rows = [] + for stage in sorted(set(baseline_stages) | set(candidate_stages)): + baseline_stage = baseline_stages.get(stage) + candidate_stage = candidate_stages.get(stage) + stage_delta_s, stage_delta_pct = _delta(baseline_stage, candidate_stage) + stage_rows.append( + { + "stage": stage, + "baseline_duration_s": baseline_stage, + "candidate_duration_s": candidate_stage, + "delta_s": stage_delta_s, + "delta_pct": stage_delta_pct, + } + ) + rows.append( + { + "case": case, + "baseline_completed": len(baseline_completed), + "candidate_completed": len(candidate_completed), + "baseline_total": len(baseline_all), + "candidate_total": len(candidate_all), + "baseline_wall_time_s": baseline_wall, + "candidate_wall_time_s": candidate_wall, + "delta_s": wall_delta_s, + "delta_pct": wall_delta_pct, + "stages": stage_rows, + } + ) + return { + "baseline_run_token": baseline.get("run_token"), + "candidate_run_token": candidate.get("run_token"), + "baseline_git_ref": baseline.get("git_ref"), + "candidate_git_ref": candidate.get("git_ref"), + "cases": rows, + } + + +def _format_number(value: Any, *, suffix: str = "") -> str: + if not isinstance(value, (int, float)): + return "-" + return f"{value:.2f}{suffix}" + + +def _print_table(rows: Sequence[Sequence[str]]) -> None: + widths = [max(len(row[index]) for row in rows) for index in range(len(rows[0]))] + for index, row in enumerate(rows): + print( + " ".join( + value.ljust(widths[column_index]) + for column_index, value in enumerate(row) + ) + ) + if index == 0: + print(" ".join("-" * width for width in widths)) + + +def _print_human(comparison: Mapping[str, Any]) -> None: + print( + f"Baseline: {comparison.get('baseline_run_token')} {comparison.get('baseline_git_ref')}" + ) + print( + f"Candidate: {comparison.get('candidate_run_token')} {comparison.get('candidate_git_ref')}" + ) + print() + rows = [ + ( + "case", + "runs", + "baseline_s", + "candidate_s", + "delta_s", + "delta_pct", + ) + ] + for case in comparison["cases"]: + rows.append( + ( + str(case["case"]), + f"{case['baseline_completed']}/{case['baseline_total']} -> " + f"{case['candidate_completed']}/{case['candidate_total']}", + _format_number(case["baseline_wall_time_s"]), + _format_number(case["candidate_wall_time_s"]), + _format_number(case["delta_s"]), + _format_number(case["delta_pct"], suffix="%"), + ) + ) + _print_table(rows) + + for case in comparison["cases"]: + stages = case["stages"] + if not stages: + continue + print() + print(f"{case['case']} stages") + stage_rows = [("stage", "baseline_s", "candidate_s", "delta_s", "delta_pct")] + for stage in stages: + stage_rows.append( + ( + str(stage["stage"]), + _format_number(stage["baseline_duration_s"]), + _format_number(stage["candidate_duration_s"]), + _format_number(stage["delta_s"]), + _format_number(stage["delta_pct"], suffix="%"), + ) + ) + _print_table(stage_rows) + + +def main() -> int: + args = _parse_args() + comparison = _comparison( + _load_summary(args.baseline), _load_summary(args.candidate) + ) + if args.json: + print(json.dumps(comparison, indent=2, sort_keys=True)) + else: + _print_human(comparison) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 617dbf291aa68dcdd4e31a3c979999c80af390bf Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 23:23:04 +0200 Subject: [PATCH 15/65] Add Rerun benchmark AWS secret refresher --- benchmark/rerun/README.md | 11 ++ benchmark/rerun/refresh_aws_secrets.py | 153 +++++++++++++++++++++++++ 2 files changed, 164 insertions(+) create mode 100644 benchmark/rerun/refresh_aws_secrets.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 4ef36903..6fd9c18a 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -7,6 +7,8 @@ This folder contains cloud benchmark harnesses for Rerun RRD workloads. job ids, stage timings, metrics, and output inspection. - `compare_results.py`: compares two benchmark `summary.json` artifacts and prints case-level and stage-level timing deltas. +- `refresh_aws_secrets.py`: copies short-lived credentials from an AWS CLI + profile into the Macrodata workspace secret environment used by cloud jobs. The default inputs are the ten base RRD files from: @@ -35,6 +37,15 @@ look better. - Workspace secrets in the selected environment must include AWS credentials for the source/output S3 bucket. The default environment is `researcher`. +If local AWS SSO credentials are valid, refresh the cloud secret environment +without printing credential values: + +```bash +uv run python benchmark/rerun/refresh_aws_secrets.py \ + --aws-profile default \ + --secret-env researcher +``` + ## Run ```bash diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py new file mode 100644 index 00000000..8deda29b --- /dev/null +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -0,0 +1,153 @@ +from __future__ import annotations + +import argparse +import json +import os +import subprocess +from collections.abc import Mapping +from typing import Any + + +DEFAULT_S3_CHECK = ( + "s3://macrodata-rerun-format-tests/dominique-sample/episode-5__base.rrd" +) +SECRET_NAMES = ( + "AWS_ACCESS_KEY_ID", + "AWS_SECRET_ACCESS_KEY", + "AWS_SESSION_TOKEN", + "AWS_DEFAULT_REGION", +) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Export short-lived AWS profile credentials into a Macrodata " + "workspace secret environment for Rerun cloud benchmarks." + ) + ) + parser.add_argument("--aws-profile", default="default") + parser.add_argument("--secret-env", default="researcher") + parser.add_argument( + "--region", + help=( + "AWS region to store as AWS_DEFAULT_REGION. Defaults to the profile " + "region, AWS_DEFAULT_REGION, AWS_REGION, or us-east-1." + ), + ) + parser.add_argument( + "--s3-check", + default=DEFAULT_S3_CHECK, + help="S3 URI to verify with the selected AWS profile before updating secrets.", + ) + parser.add_argument( + "--skip-s3-check", + action="store_true", + help="Skip the local S3 access check before writing workspace secrets.", + ) + return parser.parse_args() + + +def _run( + args: list[str], *, input_text: str | None = None +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + args, + input=input_text, + text=True, + capture_output=True, + check=True, + ) + + +def _aws_profile_arg(profile: str) -> list[str]: + return ["--profile", profile] + + +def _profile_region(profile: str) -> str | None: + try: + result = _run(["aws", "configure", "get", "region", *_aws_profile_arg(profile)]) + except subprocess.CalledProcessError: + return None + region = result.stdout.strip() + return region or None + + +def _region(args: argparse.Namespace) -> str: + return ( + args.region + or _profile_region(args.aws_profile) + or os.environ.get("AWS_DEFAULT_REGION") + or os.environ.get("AWS_REGION") + or "us-east-1" + ) + + +def _export_credentials(profile: str) -> dict[str, str]: + result = _run( + [ + "aws", + "configure", + "export-credentials", + *_aws_profile_arg(profile), + "--format", + "json", + ] + ) + payload = json.loads(result.stdout) + if not isinstance(payload, dict): + raise ValueError("aws export-credentials did not return a JSON object") + mapping = { + "AWS_ACCESS_KEY_ID": payload.get("AccessKeyId"), + "AWS_SECRET_ACCESS_KEY": payload.get("SecretAccessKey"), + "AWS_SESSION_TOKEN": payload.get("SessionToken"), + } + missing = [key for key, value in mapping.items() if not isinstance(value, str)] + if missing: + raise ValueError( + "aws export-credentials did not return required keys: " + ", ".join(missing) + ) + return {key: str(value) for key, value in mapping.items()} + + +def _check_aws_access(args: argparse.Namespace) -> None: + _run(["aws", "sts", "get-caller-identity", *_aws_profile_arg(args.aws_profile)]) + if not args.skip_s3_check: + _run(["aws", "s3", "ls", args.s3_check, *_aws_profile_arg(args.aws_profile)]) + + +def _set_secret(*, env: str, name: str, value: str) -> None: + _run( + ["macrodata", "secrets", "set", name, "--env", env, "--value-stdin"], + input_text=value, + ) + + +def _secret_payload(args: argparse.Namespace) -> dict[str, str]: + payload = _export_credentials(args.aws_profile) + payload["AWS_DEFAULT_REGION"] = _region(args) + return payload + + +def _redacted_summary(payload: Mapping[str, Any]) -> str: + return ", ".join( + f"{name}={'set' if isinstance(payload.get(name), str) and payload.get(name) else 'missing'}" + for name in SECRET_NAMES + ) + + +def main() -> int: + args = _parse_args() + _check_aws_access(args) + payload = _secret_payload(args) + for name in SECRET_NAMES: + _set_secret(env=args.secret_env, name=name, value=payload[name]) + print( + f"Updated Macrodata secrets in env {args.secret_env!r}: " + f"{_redacted_summary(payload)}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 724105460d41f9aa4df1b850050f0a00aa652ab9 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 23:25:38 +0200 Subject: [PATCH 16/65] Avoid logging Rerun benchmark secret payloads --- benchmark/rerun/refresh_aws_secrets.py | 15 ++------------- 1 file changed, 2 insertions(+), 13 deletions(-) diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py index 8deda29b..1d767ff4 100644 --- a/benchmark/rerun/refresh_aws_secrets.py +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -4,8 +4,6 @@ import json import os import subprocess -from collections.abc import Mapping -from typing import Any DEFAULT_S3_CHECK = ( @@ -129,23 +127,14 @@ def _secret_payload(args: argparse.Namespace) -> dict[str, str]: return payload -def _redacted_summary(payload: Mapping[str, Any]) -> str: - return ", ".join( - f"{name}={'set' if isinstance(payload.get(name), str) and payload.get(name) else 'missing'}" - for name in SECRET_NAMES - ) - - def main() -> int: args = _parse_args() _check_aws_access(args) payload = _secret_payload(args) for name in SECRET_NAMES: _set_secret(env=args.secret_env, name=name, value=payload[name]) - print( - f"Updated Macrodata secrets in env {args.secret_env!r}: " - f"{_redacted_summary(payload)}" - ) + updated = ", ".join(f"{name}=set" for name in SECRET_NAMES) + print(f"Updated Macrodata secrets in env {args.secret_env!r}: {updated}") return 0 From 666727aaa09d67f59a76107e8594fb0073cf6229 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Sun, 14 Jun 2026 23:29:31 +0200 Subject: [PATCH 17/65] Silence Rerun benchmark secret refresh output --- benchmark/rerun/refresh_aws_secrets.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py index 1d767ff4..b8b78e52 100644 --- a/benchmark/rerun/refresh_aws_secrets.py +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -133,8 +133,6 @@ def main() -> int: payload = _secret_payload(args) for name in SECRET_NAMES: _set_secret(env=args.secret_env, name=name, value=payload[name]) - updated = ", ".join(f"{name}=set" for name in SECRET_NAMES) - print(f"Updated Macrodata secrets in env {args.secret_env!r}: {updated}") return 0 From 234006f749f6cd4549b05fc35f7a6262e7904106 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 00:49:20 +0200 Subject: [PATCH 18/65] Use benchmark AWS profile by default --- benchmark/rerun/README.md | 6 +++--- benchmark/rerun/refresh_aws_secrets.py | 3 ++- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 6fd9c18a..f2914f29 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -37,12 +37,12 @@ look better. - Workspace secrets in the selected environment must include AWS credentials for the source/output S3 bucket. The default environment is `researcher`. -If local AWS SSO credentials are valid, refresh the cloud secret environment -without printing credential values: +If local AWS credentials are valid, refresh the cloud secret environment without +printing credential values: ```bash uv run python benchmark/rerun/refresh_aws_secrets.py \ - --aws-profile default \ + --aws-profile 210049840512_Researcher \ --secret-env researcher ``` diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py index b8b78e52..1bc3483d 100644 --- a/benchmark/rerun/refresh_aws_secrets.py +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -9,6 +9,7 @@ DEFAULT_S3_CHECK = ( "s3://macrodata-rerun-format-tests/dominique-sample/episode-5__base.rrd" ) +DEFAULT_AWS_PROFILE = "210049840512_Researcher" SECRET_NAMES = ( "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", @@ -24,7 +25,7 @@ def _parse_args() -> argparse.Namespace: "workspace secret environment for Rerun cloud benchmarks." ) ) - parser.add_argument("--aws-profile", default="default") + parser.add_argument("--aws-profile", default=DEFAULT_AWS_PROFILE) parser.add_argument("--secret-env", default="researcher") parser.add_argument( "--region", From 1bdad4ea981d88c55637fb8bcc93337ba935616b Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 00:58:38 +0200 Subject: [PATCH 19/65] Use supported AWS credential export format --- benchmark/rerun/refresh_aws_secrets.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py index 1bc3483d..a6254bcd 100644 --- a/benchmark/rerun/refresh_aws_secrets.py +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -90,7 +90,7 @@ def _export_credentials(profile: str) -> dict[str, str]: "export-credentials", *_aws_profile_arg(profile), "--format", - "json", + "process", ] ) payload = json.loads(result.stdout) From 337cafa4d7f945174cde6d001e4020e6ab1a0115 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:03:33 +0200 Subject: [PATCH 20/65] Reuse Rerun schema component maps --- src/refiner/pipeline/sources/readers/rerun.py | 34 +++++++++---------- 1 file changed, 17 insertions(+), 17 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index c41c1570..0bd0346c 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -354,10 +354,8 @@ def _robotics_row( if self.robot_type is not None: row["robot_type"] = self.robot_type - scalar_columns = _component_columns( - table=table, - component="Scalars:scalars", - ) + component_columns = _component_column_maps(table) + scalar_columns = component_columns.get("Scalars:scalars", {}) action_columns = ( _selected_columns( scalar_columns, @@ -388,10 +386,11 @@ def _robotics_row( ) row["frames"] = Tabular(frames) + image_columns = component_columns.get("EncodedImage:blob", {}) camera_columns = ( - _selected_camera_columns(table, self.videos) + _selected_camera_columns(image_columns, self.videos) if self.videos_explicit - else _camera_columns(table, self.camera_prefix) + else _camera_columns(image_columns, self.camera_prefix) ) for name, column in camera_columns.items(): values = table.column(column).combine_chunks() @@ -490,15 +489,20 @@ def _component_columns( *, component: str, ) -> dict[str, str]: - out: dict[str, str] = {} + return _component_column_maps(table).get(component, {}) + + +def _component_column_maps(table: pa.Table) -> dict[str, dict[str, str]]: + by_component: dict[str, dict[str, str]] = {} for field in table.schema: metadata = field.metadata or {} - if _metadata_text(metadata, _RERUN_COMPONENT_METADATA) != component: + field_component = _metadata_text(metadata, _RERUN_COMPONENT_METADATA) + if field_component is None: continue entity_path = _metadata_text(metadata, _RERUN_ENTITY_PATH_METADATA) if entity_path is not None: - out[entity_path] = field.name - return out + by_component.setdefault(field_component, {})[entity_path] = field.name + return by_component def _prefixed_columns(columns: Mapping[str, str], prefix: str) -> list[tuple[str, str]]: @@ -533,12 +537,9 @@ def _matches_entity_prefix(entity_path: str, prefix: str) -> bool: return entity_path == prefix or entity_path.startswith(f"{prefix}/") -def _camera_columns(table: pa.Table, prefix: str) -> dict[str, str]: +def _camera_columns(columns: Mapping[str, str], prefix: str) -> dict[str, str]: out: dict[str, str] = {} - for entity_path, column in _component_columns( - table, - component="EncodedImage:blob", - ).items(): + for entity_path, column in columns.items(): if not _matches_entity_prefix(entity_path, prefix): continue out[entity_path.strip("/").replace("/", ".")] = column @@ -546,10 +547,9 @@ def _camera_columns(table: pa.Table, prefix: str) -> dict[str, str]: def _selected_camera_columns( - table: pa.Table, + by_entity_path: Mapping[str, str], selected: Mapping[str, str], ) -> dict[str, str]: - by_entity_path = _component_columns(table, component="EncodedImage:blob") out: dict[str, str] = {} for name, path in selected.items(): column = by_entity_path.get(path) From 6b70f5903c1ac48ee57f435220294bf661b61404 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:05:47 +0200 Subject: [PATCH 21/65] Fill Rerun scalar matrices in place --- src/refiner/pipeline/sources/readers/rerun.py | 22 +++++++++---------- 1 file changed, 11 insertions(+), 11 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 0bd0346c..c899b82b 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -570,12 +570,12 @@ def _singleton_scalar_matrix( table: pa.Table, columns: Sequence[tuple[str, str]], ) -> np.ndarray: - values = [] - for _, column in columns: - values.append(_singleton_list_array(table.column(column).combine_chunks())) - if not values: - return np.empty((table.num_rows, 0), dtype=np.float64) - return np.stack(values, axis=1) + values = np.full((table.num_rows, len(columns)), np.nan, dtype=np.float64) + for index, (_, column) in enumerate(columns): + _fill_singleton_list_array( + table.column(column).combine_chunks(), values[:, index] + ) + return values def _list_column(values: np.ndarray) -> pa.Array: @@ -595,10 +595,11 @@ def _list_column(values: np.ndarray) -> pa.Array: return pa.ListArray.from_arrays(offsets, flat_values) -def _singleton_list_array(array: pa.Array) -> np.ndarray: - out = np.full(len(array), np.nan, dtype=np.float64) +def _fill_singleton_list_array(array: pa.Array, out: np.ndarray) -> None: + if len(array) != len(out): + raise ValueError("Rerun vector column length mismatch") if len(array) == 0: - return out + return if not pa.types.is_list(array.type) and not pa.types.is_large_list(array.type): raise TypeError(f"Expected a Rerun list component column, got {array.type}") offsets = np.asarray(array.offsets) @@ -606,10 +607,9 @@ def _singleton_list_array(array: pa.Array) -> np.ndarray: ends = offsets[1:] valid = np.asarray(_is_valid(array), dtype=bool) & (ends > starts) if not valid.any(): - return out + return values = np.asarray(array.values) out[valid] = values[starts[valid]] - return out def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: From e0df9ecafb1d6e7bad097587414e70c651d078bc Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:08:23 +0200 Subject: [PATCH 22/65] Cache Rerun encoded image offsets --- src/refiner/pipeline/sources/readers/rerun.py | 57 ++++++++++++------- 1 file changed, 35 insertions(+), 22 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index c899b82b..e66c477f 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -617,37 +617,50 @@ def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: from PIL import Image + encoded = _EncodedImageColumn(values) for index in range(len(values)): - data = _encoded_image_bytes(values, index) + data = encoded.bytes_at(index) if data is None: continue with Image.open(BytesIO(data)) as image: yield np.asarray(image.convert("RGB"), dtype=np.uint8) -def _encoded_image_bytes(values: pa.Array, index: int) -> bytes | None: - if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): - raise TypeError( - f"Expected a Rerun encoded image list column, got {values.type}" +class _EncodedImageColumn: + def __init__(self, values: pa.Array) -> None: + if not pa.types.is_list(values.type) and not pa.types.is_large_list( + values.type + ): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" + ) + self.values = values + self.outer_offsets = np.asarray(values.offsets) + self.valid = np.asarray(_is_valid(values), dtype=bool) + self.inner = values.values + self.inner_offsets = ( + np.asarray(self.inner.offsets) + if pa.types.is_list(self.inner.type) + or pa.types.is_large_list(self.inner.type) + else None ) - if not values[index].is_valid: - return None - outer_offsets = np.asarray(values.offsets) - outer_start = int(outer_offsets[index]) - outer_end = int(outer_offsets[index + 1]) - if outer_end <= outer_start: - return None - inner = values.values - if not pa.types.is_list(inner.type) and not pa.types.is_large_list(inner.type): - value = values[index].as_py() - if not value: + + def bytes_at(self, index: int) -> bytes | None: + if not self.valid[index]: + return None + outer_start = int(self.outer_offsets[index]) + outer_end = int(self.outer_offsets[index + 1]) + if outer_end <= outer_start: return None - return bytes(cast(bytes | bytearray | list[int], value[0])) - inner_offsets = np.asarray(inner.offsets) - byte_start = int(inner_offsets[outer_start]) - byte_end = int(inner_offsets[outer_start + 1]) - payload = inner.values.slice(byte_start, byte_end - byte_start) - return np.asarray(payload).tobytes() + if self.inner_offsets is None: + value = self.values[index].as_py() + if not value: + return None + return bytes(cast(bytes | bytearray | list[int], value[0])) + byte_start = int(self.inner_offsets[outer_start]) + byte_end = int(self.inner_offsets[outer_start + 1]) + payload = self.inner.values.slice(byte_start, byte_end - byte_start) + return np.asarray(payload).tobytes() def _is_valid(values: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: From 0f0cd224c16b261dc896b41eaa6f9e1bcb566dbd Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:19:05 +0200 Subject: [PATCH 23/65] Revert "Cache Rerun encoded image offsets" This reverts commit e0df9ecafb1d6e7bad097587414e70c651d078bc. --- src/refiner/pipeline/sources/readers/rerun.py | 57 +++++++------------ 1 file changed, 22 insertions(+), 35 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index e66c477f..c899b82b 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -617,50 +617,37 @@ def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: from PIL import Image - encoded = _EncodedImageColumn(values) for index in range(len(values)): - data = encoded.bytes_at(index) + data = _encoded_image_bytes(values, index) if data is None: continue with Image.open(BytesIO(data)) as image: yield np.asarray(image.convert("RGB"), dtype=np.uint8) -class _EncodedImageColumn: - def __init__(self, values: pa.Array) -> None: - if not pa.types.is_list(values.type) and not pa.types.is_large_list( - values.type - ): - raise TypeError( - f"Expected a Rerun encoded image list column, got {values.type}" - ) - self.values = values - self.outer_offsets = np.asarray(values.offsets) - self.valid = np.asarray(_is_valid(values), dtype=bool) - self.inner = values.values - self.inner_offsets = ( - np.asarray(self.inner.offsets) - if pa.types.is_list(self.inner.type) - or pa.types.is_large_list(self.inner.type) - else None +def _encoded_image_bytes(values: pa.Array, index: int) -> bytes | None: + if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" ) - - def bytes_at(self, index: int) -> bytes | None: - if not self.valid[index]: - return None - outer_start = int(self.outer_offsets[index]) - outer_end = int(self.outer_offsets[index + 1]) - if outer_end <= outer_start: + if not values[index].is_valid: + return None + outer_offsets = np.asarray(values.offsets) + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + return None + inner = values.values + if not pa.types.is_list(inner.type) and not pa.types.is_large_list(inner.type): + value = values[index].as_py() + if not value: return None - if self.inner_offsets is None: - value = self.values[index].as_py() - if not value: - return None - return bytes(cast(bytes | bytearray | list[int], value[0])) - byte_start = int(self.inner_offsets[outer_start]) - byte_end = int(self.inner_offsets[outer_start + 1]) - payload = self.inner.values.slice(byte_start, byte_end - byte_start) - return np.asarray(payload).tobytes() + return bytes(cast(bytes | bytearray | list[int], value[0])) + inner_offsets = np.asarray(inner.offsets) + byte_start = int(inner_offsets[outer_start]) + byte_end = int(inner_offsets[outer_start + 1]) + payload = inner.values.slice(byte_start, byte_end - byte_start) + return np.asarray(payload).tobytes() def _is_valid(values: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: From b84ff52a82631da19189e891e1139f20180cb3c2 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:21:26 +0200 Subject: [PATCH 24/65] Track Rerun metadata and output metrics --- src/refiner/pipeline/sinks/rerun.py | 2 ++ src/refiner/pipeline/sources/readers/rerun.py | 2 ++ tests/readers/test_rerun_reader.py | 15 +++++++++++++++ 3 files changed, 19 insertions(+) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index 7d440bf3..f6edd949 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -17,6 +17,7 @@ from refiner.pipeline.sources.readers.rerun import RerunRecording from refiner.utils import check_required_dependencies from refiner.worker.context import get_active_worker_token +from refiner.worker.metrics.api import log_throughput _DEFAULT_FILENAME_TEMPLATE = "{shard_id}__w{worker_id}/{row_index}.rrd" @@ -56,6 +57,7 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: segment_id=recording.segment_id, ) self._write_recording(recording, relpath) + log_throughput("files_written", 1, shard_id=shard_id, unit="files") count += 1 return count diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index c899b82b..3364c412 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -167,6 +167,8 @@ def describe(self) -> dict[str, Any]: "actions": dict(self.actions) if self.actions_explicit else None, "states": dict(self.states) if self.states_explicit else None, "videos": dict(self.videos) if self.videos_explicit else None, + "fps": self.fps, + "robot_type": self.robot_type, } ) return description diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index ad35b3c7..c9d77d89 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -198,6 +198,21 @@ def test_read_rerun_robotics_mode_with_recording_includes_recording_payload( assert row["frames"].num_rows == 3 +def test_read_rerun_describe_includes_robotics_metadata(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + description = mdr.read_rerun( + str(rrd), + output="robotics", + fps=12.5, + robot_type="testbot", + ).source.describe() + + assert description["fps"] == 12.5 + assert description["robot_type"] == "testbot" + + def test_read_rerun_robotics_mode_respects_explicit_selections( tmp_path: Path, ) -> None: From 31b96b78706f48ea79db728a8d145053f0be4b40 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:29:52 +0200 Subject: [PATCH 25/65] Harden Rerun output column collisions --- src/refiner/pipeline/sources/readers/rerun.py | 41 +++++++++++------ tests/readers/test_rerun_reader.py | 46 +++++++++++++++++++ 2 files changed, 73 insertions(+), 14 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 3364c412..f9badca0 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -36,6 +36,7 @@ _ROBOTICS_ROW_COLUMNS = frozenset( {"episode_id", "rerun", "frames", "fps", "robot_type"} ) +_RECORDING_ROW_COLUMNS = frozenset({"episode_id", "rerun"}) @dataclass(frozen=True, slots=True) @@ -129,13 +130,16 @@ def __init__( derive_names_from_paths=False, ) self.videos = _selection_map(videos, format_name="Rerun videos") + reserved_row_columns = ( + _ROBOTICS_ROW_COLUMNS if output == "robotics" else _RECORDING_ROW_COLUMNS + ) + if file_path_column in reserved_row_columns: + raise ValueError( + f"file_path_column cannot use reserved Rerun {output} row " + f"column {file_path_column!r}" + ) if output == "robotics": - if file_path_column in _ROBOTICS_ROW_COLUMNS: - raise ValueError( - f"file_path_column cannot use reserved Rerun robotics row " - f"column {file_path_column!r}" - ) - reserved_video_names = set(_ROBOTICS_ROW_COLUMNS) + reserved_video_names = set(reserved_row_columns) if file_path_column is not None: reserved_video_names.add(file_path_column) video_collisions = set(self.videos).intersection(reserved_video_names) @@ -394,6 +398,10 @@ def _robotics_row( if self.videos_explicit else _camera_columns(image_columns, self.camera_prefix) ) + reserved_video_names = set(_ROBOTICS_ROW_COLUMNS) + if self.file_path_column is not None: + reserved_video_names.add(self.file_path_column) + _validate_video_output_names(camera_columns, reserved=reserved_video_names) for name, column in camera_columns.items(): values = table.column(column).combine_chunks() row[name] = VideoFrameSequence( @@ -486,14 +494,6 @@ def _metadata_text(metadata: Mapping[bytes, bytes], key: bytes) -> str | None: return value.decode("utf-8") if value is not None else None -def _component_columns( - table: pa.Table, - *, - component: str, -) -> dict[str, str]: - return _component_column_maps(table).get(component, {}) - - def _component_column_maps(table: pa.Table) -> dict[str, dict[str, str]]: by_component: dict[str, dict[str, str]] = {} for field in table.schema: @@ -561,6 +561,19 @@ def _selected_camera_columns( return out +def _validate_video_output_names( + selected: Mapping[str, str], + *, + reserved: set[str], +) -> None: + collisions = set(selected).intersection(reserved) + if collisions: + raise ValueError( + "Rerun video output names cannot use reserved robotics row columns: " + + ", ".join(sorted(collisions)) + ) + + def _robotics_frame_table(table: pa.Table, *, timeline: str) -> pa.Table: columns: dict[str, pa.ChunkedArray] = {} if timeline in table.column_names: diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index c9d77d89..7dbb1e62 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -103,6 +103,24 @@ def _custom_robotics_rrd(path: Path) -> None: ) +def _reserved_video_name_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_reserved_video_name_test", recording_id="episode-reserved") + rr.save(path) + out = BytesIO() + Image.new("RGB", (1, 1), color=(1, 2, 3)).save(out, format="PNG") + rr.send_columns( + "/frames", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0]))], + columns=rr.EncodedImage.columns( + blob=[out.getvalue()], + media_type=["image/png"], + ), + ) + + def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -118,6 +136,16 @@ def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: assert row["file_path"] == str(rrd) +def test_read_rerun_recording_rejects_reserved_file_path_column( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + with pytest.raises(ValueError, match="reserved Rerun recording row column 'rerun'"): + mdr.read_rerun(str(rrd), file_path_column="rerun") + + def test_read_rerun_recording_preserves_sparse_rows(tmp_path: Path) -> None: rrd = tmp_path / "sparse.rrd" _sparse_rrd(rrd) @@ -213,6 +241,24 @@ def test_read_rerun_describe_includes_robotics_metadata(tmp_path: Path) -> None: assert description["robot_type"] == "testbot" +def test_read_rerun_robotics_rejects_reserved_implicit_video_name( + tmp_path: Path, +) -> None: + rrd = tmp_path / "reserved-video.rrd" + _reserved_video_name_rrd(rrd) + + with pytest.raises( + ValueError, + match="Rerun video output names cannot use reserved robotics row columns: frames", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + camera_prefix="/", + fps=30.0, + ).take(1) + + def test_read_rerun_robotics_mode_respects_explicit_selections( tmp_path: Path, ) -> None: From 3de8e90fc6411e6ad981bae48be46c670a42b416 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:32:01 +0200 Subject: [PATCH 26/65] Skip Rerun tables for raw copy benchmarks --- benchmark/rerun/README.md | 5 +- benchmark/rerun/run_cloud_benchmark.py | 6 ++- docs/reading-data/rerun.md | 5 ++ docs/writing-data/rerun.md | 4 ++ src/refiner/pipeline/pipeline.py | 15 +++--- src/refiner/pipeline/sources/readers/rerun.py | 36 +++++++------ tests/readers/test_rerun_reader.py | 54 +++++++++++++++++++ 7 files changed, 101 insertions(+), 24 deletions(-) diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index f2914f29..3f26f292 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -22,8 +22,9 @@ The default cases are: and static tables, write JSONL. - `robotics-summary`: `read_rerun(output="robotics")` for action/state paths, summarize frame rows and vector widths, write JSONL. -- `rrd-copy`: `read_rerun(output="recording").write_rerun(...)` to exercise the - distributed RRD writer's raw chunk path. +- `rrd-copy`: `read_rerun(output="recording", materialize_tables=False)` + followed by `write_rerun(...)` to exercise the distributed RRD writer's raw + chunk path without timing unused Arrow table materialization. These cases intentionally cover both the high-fidelity recording path and the robotics convenience path. Do not remove a case just to make a performance run diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index 0fcadf38..6c1af760 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -277,7 +277,11 @@ def _build_pipeline( .write_jsonl(output) ) if case == "rrd-copy": - return mdr.read_rerun(inputs, output="recording").write_rerun(output) + return mdr.read_rerun( + inputs, + output="recording", + materialize_tables=False, + ).write_rerun(output) raise ValueError(f"Unsupported benchmark case: {case}") diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index f783c973..5b24ae47 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -50,6 +50,11 @@ Each row includes: tables returned. If `timelines` is omitted, the reader materializes all timeline indexes reported by the Rerun schema. +For raw RRD copy workflows that immediately call `write_rerun`, set +`materialize_tables=False`. The row still carries the source recording metadata +needed by the writer's chunk-copy path, but skips the Arrow timeline/static +tables that downstream code will not inspect. + ## Robotics rows With `output="robotics"`, the reader creates rows that can be passed to diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index 70fa576e..bcb30c9a 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -42,6 +42,10 @@ When the input row came from `read_rerun`, the writer uses Rerun's raw preserves Rerun chunk metadata and avoids re-emitting large Arrow tables through Python. +For pure copy jobs, use `read_rerun(..., materialize_tables=False)` before +`write_rerun(...)` to skip timeline/static table materialization while keeping +the raw source chunks available to the writer. + If a `RerunRecording` has no source file, the writer falls back to table emission with `send_dataframe`. Static Rerun component columns are sent as static data, and dynamic timeline tables are sent separately. The same fallback diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 9ea9fbb8..2764f7c9 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -1257,6 +1257,7 @@ def read_rerun( timelines: Sequence[str] | None = None, primary_timeline: str | None = None, include_static: bool = True, + materialize_tables: bool = True, include_recording: bool | None = None, fill_latest_at: bool = False, action_prefix: str = "/action", @@ -1272,12 +1273,13 @@ def read_rerun( RRD files are planned as atomic input shards. With ``output="recording"``, each emitted row preserves the selected Rerun data as Arrow-backed - ``Tabular`` tables grouped by timeline under the ``rerun`` field. With - ``output="robotics"``, the reader additionally derives common robotics - episode fields from configurable Rerun entity prefixes so the rows can be - passed through ``to_robot_rows(...)`` and robotics writers. Pass ``actions``, - ``states``, or ``videos`` to pin exact entity paths and output order instead - of using prefix-derived defaults. + ``Tabular`` tables grouped by timeline under the ``rerun`` field. Set + ``materialize_tables=False`` for raw ``write_rerun`` copy workflows that + only need source chunk metadata. With ``output="robotics"``, the reader + additionally derives common robotics episode fields from configurable Rerun + entity prefixes so the rows can be passed through ``to_robot_rows(...)`` and + robotics writers. Pass ``actions``, ``states``, or ``videos`` to pin exact + entity paths and output order instead of using prefix-derived defaults. """ return RefinerPipeline( source=RerunReader( @@ -1293,6 +1295,7 @@ def read_rerun( timelines=timelines, primary_timeline=primary_timeline, include_static=include_static, + materialize_tables=materialize_tables, include_recording=include_recording, fill_latest_at=fill_latest_at, action_prefix=action_prefix, diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index f9badca0..b35bb91f 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -76,6 +76,7 @@ def __init__( timelines: Sequence[str] | None = None, primary_timeline: str | None = None, include_static: bool = True, + materialize_tables: bool = True, include_recording: bool | None = None, fill_latest_at: bool = False, action_prefix: str = "/action", @@ -109,6 +110,7 @@ def __init__( self.timelines = tuple(timelines) if timelines is not None else None self.primary_timeline = primary_timeline self.include_static = include_static + self.materialize_tables = materialize_tables self.include_recording = ( output == "recording" if include_recording is None else include_recording ) @@ -163,6 +165,7 @@ def describe(self) -> dict[str, Any]: "timelines": self.timelines, "primary_timeline": self.primary_timeline, "include_static": self.include_static, + "materialize_tables": self.materialize_tables, "include_recording": self.include_recording, "fill_latest_at": self.fill_latest_at, "action_prefix": self.action_prefix, @@ -241,23 +244,26 @@ def _read_dataset( timelines=timelines, ) else: - content_view = self._view_for_contents(view) - static = ( - _collect_table(content_view.reader(index=None)) - if self.include_static - else None - ) - tables = { - timeline: Tabular( - _collect_table( - content_view.reader( - index=timeline, - fill_latest_at=self.fill_latest_at, + tables: dict[str, Tabular] = {} + static = None + if self.materialize_tables: + content_view = self._view_for_contents(view) + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) + tables = { + timeline: Tabular( + _collect_table( + content_view.reader( + index=timeline, + fill_latest_at=self.fill_latest_at, + ) ) ) - ) - for timeline in timelines - } + for timeline in timelines + } data: dict[str, Any] = { "episode_id": segment_id, "rerun": RerunRecording( diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 7dbb1e62..3eae2a68 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -157,6 +157,29 @@ def test_read_rerun_recording_preserves_sparse_rows(tmp_path: Path) -> None: assert table.column("frame").to_pylist() == [0, 1, 2] +def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + recording = row["rerun"] + + assert recording.recording_id == "episode-a" + assert recording.tables == {} + assert recording.static is None + assert recording.source_file is not None + assert recording.timelines == ("frame",) + + def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -381,6 +404,37 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: assert recording.tables["frame"].num_rows == 3 +def test_write_rerun_uses_source_chunks_without_materialized_tables( + tmp_path: Path, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-raw-copy" + _tiny_rrd(source) + + row = cast( + Any, + next( + mdr.read_rerun( + str(source), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", [row]) + sink.on_shard_complete("shard-a") + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + copied = cast( + Any, next(mdr.read_rerun(str(written[0]), timelines=("frame",)).source.read()) + ) + + assert copied["episode_id"] == "episode-a" + assert copied["rerun"].tables["frame"].num_rows == 3 + + def test_write_rerun_without_footer_uses_table_fallback( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From 7c6856acacfb7716011e52844a6261561c723578 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:44:08 +0200 Subject: [PATCH 27/65] Optimize Rerun raw copy path --- src/refiner/pipeline/sinks/rerun.py | 113 +++++++++------- src/refiner/pipeline/sources/readers/rerun.py | 126 +++++++++++++++--- tests/readers/test_rerun_reader.py | 99 ++++++++++++++ 3 files changed, 274 insertions(+), 64 deletions(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index f6edd949..bb94845c 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -14,7 +14,7 @@ from refiner.pipeline.data.row import Row from refiner.pipeline.sinks.base import BaseSink from refiner.pipeline.sinks.reducer.file import FileCleanupReducerSink -from refiner.pipeline.sources.readers.rerun import RerunRecording +from refiner.pipeline.sources.readers.rerun import RerunRecording, _local_rrd from refiner.utils import check_required_dependencies from refiner.worker.context import get_active_worker_token from refiner.worker.metrics.api import log_throughput @@ -131,32 +131,56 @@ def _write_source_chunks( *, application_id: str, ) -> None: - import rerun as rr - source = recording.source_file if source is None: raise ValueError("Rerun source chunk write requires source_file") - with _local_rrd(source) as local_path: - reader = rr.experimental.RrdReader(local_path) - store = _matching_store(reader, recording) - stream = reader.stream(store=store) - if recording.contents is not None: - stream = stream.filter(content=recording.contents) - if not recording.include_static: - stream = stream.drop(is_static=True) - stream = _filter_timelines( - stream, - reader=reader, - store=store, - recording=recording, + local_source_path = recording.local_source_path + if local_source_path is not None and local_source_path.exists(): + _write_source_chunks_from_path( + recording, + path, + local_path=local_source_path, + application_id=application_id, ) - stream.write_rrd( + return + with _local_rrd(source) as local_path: + _write_source_chunks_from_path( + recording, path, - application_id=recording.application_id or application_id, - recording_id=recording.recording_id or recording.segment_id, + local_path=local_path, + application_id=application_id, ) +def _write_source_chunks_from_path( + recording: RerunRecording, + path: Path, + *, + local_path: Path, + application_id: str, +) -> None: + import rerun as rr + + reader = rr.experimental.RrdReader(local_path) + store = _matching_store(reader, recording) + stream = reader.stream(store=store) + if recording.contents is not None: + stream = stream.filter(content=recording.contents) + if not recording.include_static: + stream = stream.drop(is_static=True) + stream = _filter_timelines( + stream, + reader=reader, + store=store, + recording=recording, + ) + stream.write_rrd( + path, + application_id=recording.application_id or application_id, + recording_id=recording.recording_id or recording.segment_id, + ) + + def _filter_timelines( stream: Any, *, @@ -259,26 +283,6 @@ def _is_static_column(field: pa.Field) -> bool: return (field.metadata or {}).get(b"rerun:is_static") == b"true" -class _local_rrd: - def __init__(self, source: DataFile) -> None: - self.source = source - self.tmpdir: tempfile.TemporaryDirectory[str] | None = None - self.path: Path | None = None - - def __enter__(self) -> Path: - if self.source.is_local: - return Path(self.source.abs_path()) - self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-source-") - name = os.path.basename(self.source.path) or "recording.rrd" - self.path = Path(self.tmpdir.name) / name - self.source.copy(str(self.path)) - return self.path - - def __exit__(self, *args: object) -> None: - if self.tmpdir is not None: - self.tmpdir.cleanup() - - def _validate_filename_template(filename_template: str) -> None: fields: set[str] = set() for _literal_text, field_name, format_spec, conversion in Formatter().parse( @@ -319,17 +323,36 @@ def _render_relpath( row_index: int, segment_id: str, ) -> str: + field_values: dict[str, object] = { + "shard_id": shard_id, + "worker_id": worker_id, + "row_index": row_index, + "segment_id": segment_id, + } + if "segment_id" in _template_fields(filename_template): + field_values["segment_id"] = _normalize_path_segment(segment_id, "segment_id") return _normalize_relpath( - filename_template.format( - shard_id=shard_id, - worker_id=worker_id, - row_index=row_index, - segment_id=segment_id, - ), + filename_template.format(**field_values), "rendered filename", ) +def _template_fields(filename_template: str) -> set[str]: + return { + field_name + for _literal_text, field_name, _format_spec, _conversion in Formatter().parse( + filename_template + ) + if field_name is not None + } + + +def _normalize_path_segment(value: str, label: str) -> str: + if not value or value in {".", ".."} or "/" in value or "\\" in value: + raise ValueError(f"{label} must be a single relative path segment") + return value + + def _normalize_relpath(path: str, label: str) -> str: if path.startswith("/"): raise ValueError(f"{label} must be relative") diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index b35bb91f..8cab213e 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -48,6 +48,7 @@ class RerunRecording: tables: Mapping[str, Tabular] static: Tabular | None = None source_file: DataFile | None = None + local_source_path: Path | None = None application_id: str | None = None recording_id: str | None = None contents: tuple[str, ...] | None = None @@ -193,6 +194,29 @@ def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: def _read_files( self, local_files: Sequence[tuple[DataFile, Path]], + ) -> Iterator[SourceUnit]: + if self.output == "recording" and not self.materialize_tables: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + server_fallback: list[tuple[DataFile, Path]] = [] + for source, local_path in local_files: + rows = self._read_metadata_only_recording_rows(source, local_path) + if rows: + yield from rows + else: + server_fallback.append((source, local_path)) + if server_fallback: + yield from self._read_files_with_server(server_fallback) + return + + yield from self._read_files_with_server(local_files) + + def _read_files_with_server( + self, + local_files: Sequence[tuple[DataFile, Path]], ) -> Iterator[SourceUnit]: check_required_dependencies( "read_rerun", @@ -239,6 +263,7 @@ def _read_dataset( segment_id=segment_id, source_path=source.abs_path(), source_file=source, + local_source_path=local_path, application_id=application_id, recording_id=recording_id, timelines=timelines, @@ -264,23 +289,66 @@ def _read_dataset( ) for timeline in timelines } - data: dict[str, Any] = { - "episode_id": segment_id, - "rerun": RerunRecording( - segment_id=segment_id, - source_path=source.abs_path(), - tables=tables, - static=Tabular(static) if static is not None else None, - source_file=source, - application_id=application_id, - recording_id=recording_id, - contents=self.contents, - timelines=self.timelines, - include_static=self.include_static, - ), - } - self._with_file_path(data, source) - yield DictRow(data) + yield self._recording_row( + segment_id=segment_id, + source=source, + local_path=local_path, + tables=tables, + static=Tabular(static) if static is not None else None, + application_id=application_id, + recording_id=recording_id, + ) + + def _read_metadata_only_recording_rows( + self, + source: DataFile, + local_path: Path, + ) -> list[DictRow]: + rows = [] + for store in _recording_entries(local_path): + recording_id = str(store.recording_id) + rows.append( + self._recording_row( + segment_id=recording_id, + source=source, + local_path=local_path, + tables={}, + static=None, + application_id=store.application_id, + recording_id=recording_id, + ) + ) + return rows + + def _recording_row( + self, + *, + segment_id: str, + source: DataFile, + local_path: Path, + tables: Mapping[str, Tabular], + static: Tabular | None, + application_id: str | None, + recording_id: str | None, + ) -> DictRow: + data: dict[str, Any] = { + "episode_id": segment_id, + "rerun": RerunRecording( + segment_id=segment_id, + source_path=source.abs_path(), + tables=tables, + static=static, + source_file=source, + local_source_path=local_path, + application_id=application_id, + recording_id=recording_id, + contents=self.contents, + timelines=self.timelines, + include_static=self.include_static, + ), + } + self._with_file_path(data, source) + return DictRow(data) def _timelines(self, schema: Any) -> tuple[str, ...]: if self.timelines is not None: @@ -327,14 +395,16 @@ def _robotics_row( segment_id: str, source_path: str, source_file: DataFile, + local_source_path: Path, application_id: str | None, recording_id: str, timelines: Sequence[str], ) -> DictRow: timeline = self._primary_timeline(timelines) contents = self._robotics_contents() + content_view = view.filter_contents(contents) table = _collect_table( - view.filter_contents(contents).reader( + content_view.reader( index=timeline, fill_latest_at=self.fill_latest_at, ) @@ -345,7 +415,7 @@ def _robotics_row( } if self.include_recording: static = ( - _collect_table(view.filter_contents(contents).reader(index=None)) + _collect_table(content_view.reader(index=None)) if self.include_static else None ) @@ -355,6 +425,7 @@ def _robotics_row( tables={timeline: Tabular(table)}, static=Tabular(static) if static is not None else None, source_file=source_file, + local_source_path=local_source_path, application_id=application_id, recording_id=recording_id, contents=tuple(contents), @@ -410,6 +481,7 @@ def _robotics_row( _validate_video_output_names(camera_columns, reserved=reserved_video_names) for name, column in camera_columns.items(): values = table.column(column).combine_chunks() + _require_dense_encoded_images(values, video_name=name) row[name] = VideoFrameSequence( lambda values=values: _iter_encoded_images(values), fps=self.fps or 30.0, @@ -633,6 +705,22 @@ def _fill_singleton_list_array(array: pa.Array, out: np.ndarray) -> None: out[valid] = values[starts[valid]] +def _require_dense_encoded_images(values: pa.Array, *, video_name: str) -> None: + if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): + raise TypeError( + f"Expected a Rerun encoded image list column, got {values.type}" + ) + offsets = np.asarray(values.offsets) + missing = np.asarray(_is_valid(values), dtype=bool) == 0 + if len(offsets) > 1: + missing |= offsets[1:] <= offsets[:-1] + if missing.any(): + raise ValueError( + f"Rerun video {video_name!r} has missing frames on the primary timeline; " + "use fill_latest_at=True or select a denser timeline" + ) + + def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: from io import BytesIO diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 3eae2a68..621dce4c 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -9,8 +9,10 @@ import refiner as mdr from refiner.pipeline import Row +from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table +from refiner.pipeline.sources.readers.rerun import RerunRecording pytest.importorskip("rerun") @@ -103,6 +105,29 @@ def _custom_robotics_rrd(path: Path) -> None: ) +def _sparse_camera_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_sparse_camera_test", recording_id="episode-sparse-camera") + rr.save(path) + rr.send_columns( + "/robot/actions/x", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0, 1]))], + columns=rr.Scalars.columns(scalars=np.asarray([1.0, 2.0])), + ) + out = BytesIO() + Image.new("RGB", (1, 1), color=(1, 2, 3)).save(out, format="PNG") + rr.send_columns( + "/robot/cameras/top", + indexes=[rr.TimeColumn("frame", sequence=np.asarray([0]))], + columns=rr.EncodedImage.columns( + blob=[out.getvalue()], + media_type=["image/png"], + ), + ) + + def _reserved_video_name_rrd(path: Path) -> None: import rerun as rr from PIL import Image @@ -177,9 +202,39 @@ def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> assert recording.tables == {} assert recording.static is None assert recording.source_file is not None + assert recording.local_source_path == rrd assert recording.timelines == ("frame",) +def test_read_rerun_recording_without_materialized_tables_skips_server( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun.RerunReader._read_files_with_server", + lambda *args, **kwargs: pytest.fail( + "metadata-only recording rows do not need the Rerun server" + ), + ) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + + assert row["episode_id"] == "episode-a" + assert row["rerun"].tables == {} + + def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -316,6 +371,26 @@ def test_read_rerun_robotics_mode_respects_explicit_selections( ] +def test_read_rerun_robotics_mode_rejects_sparse_video_stream( + tmp_path: Path, +) -> None: + rrd = tmp_path / "sparse-camera.rrd" + _sparse_camera_rrd(rrd) + + with pytest.raises( + ValueError, + match="missing frames on the primary timeline", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + timelines=("frame",), + actions=("/robot/actions/x",), + videos={"observation.images.top": "/robot/cameras/top"}, + fps=5.0, + ).take(1) + + def test_read_rerun_robotics_mode_with_explicit_timeline_uses_table_metadata( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -406,6 +481,7 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: def test_write_rerun_uses_source_chunks_without_materialized_tables( tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, ) -> None: source = tmp_path / "tiny.rrd" output = tmp_path / "out-raw-copy" @@ -421,6 +497,12 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( ).source.read() ), ) + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun._local_rrd", + lambda *args, **kwargs: pytest.fail( + "writer should reuse the reader-local RRD path while it is available" + ), + ) sink = RerunSink(str(output)) sink.write_shard_block("shard-a", [row]) sink.on_shard_complete("shard-a") @@ -435,6 +517,23 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( assert copied["rerun"].tables["frame"].num_rows == 3 +def test_write_rerun_rejects_segment_id_path_separator_in_filename( + tmp_path: Path, +) -> None: + recording = RerunRecording( + segment_id="episode/5", + source_path="memory://episode.rrd", + tables={}, + ) + sink = RerunSink( + str(tmp_path / "out-segment-id"), + filename_template="{shard_id}__w{worker_id}/{segment_id}.rrd", + ) + + with pytest.raises(ValueError, match="segment_id must be a single"): + sink.write_shard_block("shard-a", [DictRow({"rerun": recording})]) + + def test_write_rerun_without_footer_uses_table_fallback( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From 54c31ad6acce25295d26b45cbf047558da45f8c4 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:47:12 +0200 Subject: [PATCH 28/65] Track Rerun benchmark shard planning --- benchmark/rerun/README.md | 4 ++++ benchmark/rerun/run_cloud_benchmark.py | 28 ++++++++++++++++++++++++++ docs/writing-data/rerun.md | 4 +++- 3 files changed, 35 insertions(+), 1 deletion(-) diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 3f26f292..c25d90b8 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -71,6 +71,10 @@ Artifacts are written under `benchmark/rerun/artifacts/` by default: - one per-case result JSON - one summary JSON for the benchmark session +Each case records `planned_shards`. RRD files are file-atomic, so runs where +`planned_shards < --num-workers` can underutilize workers and should not be used +as scaling evidence. + ## Compare After running a baseline and candidate benchmark, compare their summaries: diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index 6c1af760..788697ab 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -61,6 +61,8 @@ class CaseResult: started_at_utc: str finished_at_utc: str input_count: int + planned_shards: int | None + planning_warning: str | None output_root: str cloud_wall_time_s: float | None queue_time_s: float | None @@ -422,6 +424,24 @@ def _inspect_output(path: str) -> tuple[int | None, int | None, str | None]: return None, None, str(err) +def _planned_shard_count( + pipeline: mdr.RefinerPipeline, + *, + requested_workers: int, +) -> tuple[int | None, str | None]: + try: + planned_shards = len(pipeline.list_shards()) + except Exception as err: + return None, f"could not inspect planned shards before launch: {err}" + if planned_shards < requested_workers: + return ( + planned_shards, + "planned Rerun shards are fewer than requested workers; file-atomic " + "RRD sharding may underutilize cloud workers", + ) + return planned_shards, None + + def _run_case( *, args: argparse.Namespace, @@ -445,6 +465,12 @@ def _run_case( timeline=args.timeline, fps=args.fps, ) + planned_shards, planning_warning = _planned_shard_count( + pipeline, + requested_workers=args.num_workers, + ) + if planning_warning is not None: + print(f"Warning: {case}: {planning_warning}", file=sys.stderr, flush=True) started_at = _utc_now() os.environ.setdefault("REFINER_ATTACH", "detach") launch = pipeline.launch_cloud( @@ -476,6 +502,8 @@ def _run_case( started_at_utc=started_at, finished_at_utc=finished_at, input_count=len(inputs), + planned_shards=planned_shards, + planning_warning=planning_warning, output_root=output, cloud_wall_time_s=_duration_s(job.get("startedAt"), job.get("endedAt")), queue_time_s=_duration_s(job.get("createdAt"), job.get("startedAt")), diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index bcb30c9a..af5b4b88 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -33,7 +33,9 @@ template is: The template must include `{shard_id}` and `{worker_id}` so retry cleanup can distinguish finalized worker outputs from abandoned attempt outputs. You can -also use `{row_index}` and `{segment_id}`. +also use `{row_index}` and `{segment_id}`. When `{segment_id}` is present, the +recording segment id must be a single path segment; ids containing `/`, `\`, +`.`, or `..` are rejected before writing. ## Writer strategy From 99918bfd03ea9742ce03fc6be2fcbeef6fb1c73f Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 01:53:49 +0200 Subject: [PATCH 29/65] Harden Rerun writer edge cases --- benchmark/rerun/README.md | 2 + benchmark/rerun/compare_results.py | 56 ++++++++++++++ benchmark/rerun/run_cloud_benchmark.py | 22 +++--- docs/reading-data/rerun.md | 6 +- docs/writing-data/rerun.md | 13 ++-- src/refiner/pipeline/sinks/rerun.py | 50 ++++++++---- src/refiner/pipeline/sources/readers/rerun.py | 13 ++-- tests/readers/test_rerun_reader.py | 77 +++++++++++++++++-- 8 files changed, 192 insertions(+), 47 deletions(-) diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index c25d90b8..4ab23bc3 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -87,3 +87,5 @@ uv run python benchmark/rerun/compare_results.py \ Only completed jobs are used for timing deltas. Failed jobs still appear in the run-count columns so setup problems are visible instead of silently averaged in. +Planned shard counts and shard-planning warnings are printed with the timing +table. diff --git a/benchmark/rerun/compare_results.py b/benchmark/rerun/compare_results.py index 43019c86..7c92e443 100644 --- a/benchmark/rerun/compare_results.py +++ b/benchmark/rerun/compare_results.py @@ -56,6 +56,24 @@ def _group_results(summary: Mapping[str, Any]) -> dict[str, list[Mapping[str, An return dict(grouped) +def _unique_values(results: Sequence[Mapping[str, Any]], key: str) -> list[Any]: + values = [] + for result in results: + value = result.get(key) + if value is not None and value not in values: + values.append(value) + return values + + +def _warnings(results: Sequence[Mapping[str, Any]], key: str) -> list[str]: + values = [] + for result in results: + value = result.get(key) + if isinstance(value, str) and value and value not in values: + values.append(value) + return values + + def _stage_key(stage: Mapping[str, Any]) -> str: name = stage.get("name") if isinstance(name, str) and name: @@ -130,6 +148,18 @@ def _comparison( "candidate_completed": len(candidate_completed), "baseline_total": len(baseline_all), "candidate_total": len(candidate_all), + "baseline_planned_shards": _unique_values( + baseline_all, "planned_shards" + ), + "candidate_planned_shards": _unique_values( + candidate_all, "planned_shards" + ), + "baseline_planning_warnings": _warnings( + baseline_all, "planning_warning" + ), + "candidate_planning_warnings": _warnings( + candidate_all, "planning_warning" + ), "baseline_wall_time_s": baseline_wall, "candidate_wall_time_s": candidate_wall, "delta_s": wall_delta_s, @@ -152,6 +182,12 @@ def _format_number(value: Any, *, suffix: str = "") -> str: return f"{value:.2f}{suffix}" +def _format_values(values: Sequence[Any]) -> str: + if not values: + return "-" + return ",".join(str(value) for value in values) + + def _print_table(rows: Sequence[Sequence[str]]) -> None: widths = [max(len(row[index]) for row in rows) for index in range(len(rows[0]))] for index, row in enumerate(rows): @@ -177,6 +213,7 @@ def _print_human(comparison: Mapping[str, Any]) -> None: ( "case", "runs", + "shards", "baseline_s", "candidate_s", "delta_s", @@ -189,6 +226,8 @@ def _print_human(comparison: Mapping[str, Any]) -> None: str(case["case"]), f"{case['baseline_completed']}/{case['baseline_total']} -> " f"{case['candidate_completed']}/{case['candidate_total']}", + f"{_format_values(case['baseline_planned_shards'])} -> " + f"{_format_values(case['candidate_planned_shards'])}", _format_number(case["baseline_wall_time_s"]), _format_number(case["candidate_wall_time_s"]), _format_number(case["delta_s"]), @@ -216,6 +255,23 @@ def _print_human(comparison: Mapping[str, Any]) -> None: ) _print_table(stage_rows) + warning_rows = [("case", "side", "planned_shards", "warning")] + for case in comparison["cases"]: + for side in ("baseline", "candidate"): + for warning in case[f"{side}_planning_warnings"]: + warning_rows.append( + ( + str(case["case"]), + side, + _format_values(case[f"{side}_planned_shards"]), + warning, + ) + ) + if len(warning_rows) > 1: + print() + print("Planning warnings") + _print_table(warning_rows) + def main() -> int: args = _parse_args() diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index 788697ab..def679de 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -3,7 +3,7 @@ import argparse import json import os -import platform +import platform as platform_module import re import subprocess import sys @@ -70,10 +70,10 @@ class CaseResult: output_file_count: int | None output_size_bytes: int | None output_inspection_error: str | None - python_version: str - platform: str + submitter_python_version: str + submitter_platform: str git_ref: str - package_versions: dict[str, str] + submitter_package_versions: dict[str, str] def summarize_recording(row: Row) -> DictRow: @@ -413,8 +413,12 @@ def _inspect_output(path: str) -> tuple[int | None, int | None, str | None]: return 0, 0, None total_size = 0 total_files = 0 - for child in fs.find(fs_path): - info = fs.info(child) + found = fs.find(fs_path, detail=True) + if isinstance(found, dict): + infos = found.values() + else: + infos = (fs.info(child) for child in fs.find(fs_path)) + for info in infos: if info.get("type") == "directory": continue total_files += 1 @@ -511,10 +515,10 @@ def _run_case( output_file_count=output_file_count, output_size_bytes=output_size_bytes, output_inspection_error=output_error, - python_version=sys.version.replace("\n", " "), - platform=platform.platform(), + submitter_python_version=sys.version.replace("\n", " "), + submitter_platform=platform_module.platform(), git_ref=git_ref, - package_versions=_package_versions(), + submitter_package_versions=_package_versions(), ) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index 5b24ae47..deaca638 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -51,9 +51,9 @@ tables returned. If `timelines` is omitted, the reader materializes all timeline indexes reported by the Rerun schema. For raw RRD copy workflows that immediately call `write_rerun`, set -`materialize_tables=False`. The row still carries the source recording metadata -needed by the writer's chunk-copy path, but skips the Arrow timeline/static -tables that downstream code will not inspect. +`materialize_tables=False` with `output="recording"`. The row still carries the +source recording metadata needed by the writer's chunk-copy path, but skips the +Arrow timeline/static tables that downstream code will not inspect. ## Robotics rows diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index af5b4b88..e8874b49 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -32,10 +32,11 @@ template is: ``` The template must include `{shard_id}` and `{worker_id}` so retry cleanup can -distinguish finalized worker outputs from abandoned attempt outputs. You can -also use `{row_index}` and `{segment_id}`. When `{segment_id}` is present, the -recording segment id must be a single path segment; ids containing `/`, `\`, -`.`, or `..` are rejected before writing. +distinguish finalized worker outputs from abandoned attempt outputs. It must +also include `{row_index}` or `{segment_id}` so each input row writes a distinct +RRD file. When `{segment_id}` is present, the recording segment id must be a +single path segment; ids containing `/`, `\`, `.`, or `..` are rejected before +writing. ## Writer strategy @@ -52,7 +53,9 @@ If a `RerunRecording` has no source file, the writer falls back to table emission with `send_dataframe`. Static Rerun component columns are sent as static data, and dynamic timeline tables are sent separately. The same fallback is used when `write_footer=False`, because Rerun's raw chunk writer always -writes footer metadata. +writes footer metadata. No-footer writes require materialized Rerun table data; +metadata-only rows from `materialize_tables=False` should use the default +`write_footer=True` raw chunk path. ## Reducer diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index bb94845c..4a8e4344 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -39,6 +39,7 @@ def __init__( self.app_id = app_id self.write_footer = write_footer self._row_indices: dict[str, int] = {} + self._written_relpaths: set[str] = set() def _declared_refiner_extras(self) -> tuple[str, ...]: return ("rerun",) @@ -56,7 +57,13 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: row_index=row_index, segment_id=recording.segment_id, ) + if relpath in self._written_relpaths: + raise ValueError( + "write_rerun filename_template rendered duplicate output path " + f"{relpath!r}; include {{row_index}} or another unique row field" + ) self._write_recording(recording, relpath) + self._written_relpaths.add(relpath) log_throughput("files_written", 1, shard_id=shard_id, unit="files") count += 1 return count @@ -134,15 +141,6 @@ def _write_source_chunks( source = recording.source_file if source is None: raise ValueError("Rerun source chunk write requires source_file") - local_source_path = recording.local_source_path - if local_source_path is not None and local_source_path.exists(): - _write_source_chunks_from_path( - recording, - path, - local_path=local_source_path, - application_id=application_id, - ) - return with _local_rrd(source) as local_path: _write_source_chunks_from_path( recording, @@ -241,19 +239,32 @@ def _write_recording_tables( ) -> None: import rerun as rr + static = ( + _sendable_static_table(recording.static.table) + if recording.static is not None + else None + ) + dynamic_tables = [ + dynamic + for table in recording.tables.values() + if (dynamic := _sendable_dynamic_table(table.table)).num_columns > 0 + ] + if (static is None or static.num_columns == 0) and not dynamic_tables: + raise ValueError( + "write_rerun cannot write a RerunRecording without materialized " + "Rerun table columns; use write_footer=True for raw source chunk " + "writes or read_rerun(..., materialize_tables=True)" + ) + with rr.RecordingStream( recording.application_id or application_id, recording_id=recording.recording_id or recording.segment_id, ) as rec: rec.save(path, write_footer=write_footer) - if recording.static is not None: - static = _sendable_static_table(recording.static.table) - if static.num_columns > 0: - rec.send_dataframe(static) - for table in recording.tables.values(): - dynamic = _sendable_dynamic_table(table.table) - if dynamic.num_columns > 0: - rec.send_dataframe(dynamic) + if static is not None and static.num_columns > 0: + rec.send_dataframe(static) + for dynamic in dynamic_tables: + rec.send_dataframe(dynamic) def _sendable_static_table(table: pa.Table) -> pa.Table: @@ -304,6 +315,11 @@ def _validate_filename_template(filename_template: str) -> None: "filename_template requires fields: " + ", ".join(f"{{{field}}}" for field in sorted(missing)) ) + if not fields.intersection({"row_index", "segment_id"}): + raise ValueError( + "filename_template requires {row_index} or {segment_id} so each " + "Rerun row writes a distinct file" + ) _normalize_relpath( filename_template.format( shard_id="shard", diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 8cab213e..0998b0f0 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -48,7 +48,6 @@ class RerunRecording: tables: Mapping[str, Tabular] static: Tabular | None = None source_file: DataFile | None = None - local_source_path: Path | None = None application_id: str | None = None recording_id: str | None = None contents: tuple[str, ...] | None = None @@ -91,6 +90,14 @@ def __init__( ) -> None: if output not in ("recording", "robotics"): raise ValueError("output must be 'recording' or 'robotics'") + if output == "robotics" and not materialize_tables: + raise ValueError( + "materialize_tables=False is only supported for recording output" + ) + if output == "recording" and include_recording is False: + raise ValueError( + "include_recording=False is only supported for robotics output" + ) if fps is not None: fps = float(fps) if not np.isfinite(fps) or fps <= 0: @@ -263,7 +270,6 @@ def _read_dataset( segment_id=segment_id, source_path=source.abs_path(), source_file=source, - local_source_path=local_path, application_id=application_id, recording_id=recording_id, timelines=timelines, @@ -339,7 +345,6 @@ def _recording_row( tables=tables, static=static, source_file=source, - local_source_path=local_path, application_id=application_id, recording_id=recording_id, contents=self.contents, @@ -395,7 +400,6 @@ def _robotics_row( segment_id: str, source_path: str, source_file: DataFile, - local_source_path: Path, application_id: str | None, recording_id: str, timelines: Sequence[str], @@ -425,7 +429,6 @@ def _robotics_row( tables={timeline: Tabular(table)}, static=Tabular(static) if static is not None else None, source_file=source_file, - local_source_path=local_source_path, application_id=application_id, recording_id=recording_id, contents=tuple(contents), diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 621dce4c..b9adee6f 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -202,10 +202,26 @@ def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> assert recording.tables == {} assert recording.static is None assert recording.source_file is not None - assert recording.local_source_path == rrd assert recording.timelines == ("frame",) +def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + with pytest.raises( + ValueError, + match="materialize_tables=False is only supported for recording output", + ): + mdr.read_rerun(str(rrd), output="robotics", materialize_tables=False) + + with pytest.raises( + ValueError, + match="include_recording=False is only supported for robotics output", + ): + mdr.read_rerun(str(rrd), output="recording", include_recording=False) + + def test_read_rerun_recording_without_materialized_tables_skips_server( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -481,7 +497,6 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: def test_write_rerun_uses_source_chunks_without_materialized_tables( tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, ) -> None: source = tmp_path / "tiny.rrd" output = tmp_path / "out-raw-copy" @@ -497,12 +512,6 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( ).source.read() ), ) - monkeypatch.setattr( - "refiner.pipeline.sinks.rerun._local_rrd", - lambda *args, **kwargs: pytest.fail( - "writer should reuse the reader-local RRD path while it is available" - ), - ) sink = RerunSink(str(output)) sink.write_shard_block("shard-a", [row]) sink.on_shard_complete("shard-a") @@ -534,6 +543,35 @@ def test_write_rerun_rejects_segment_id_path_separator_in_filename( sink.write_shard_block("shard-a", [DictRow({"rerun": recording})]) +def test_write_rerun_rejects_non_row_varying_filename_template( + tmp_path: Path, +) -> None: + with pytest.raises( + ValueError, match="requires \\{row_index\\} or \\{segment_id\\}" + ): + RerunSink( + str(tmp_path / "out-overwrite"), + filename_template="{shard_id}__w{worker_id}.rrd", + ) + + +def test_write_rerun_rejects_duplicate_rendered_filename( + tmp_path: Path, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-duplicate" + _tiny_rrd(source) + unit = next(mdr.read_rerun(str(source), timelines=("frame",)).source.read()) + assert isinstance(unit, Row) + + sink = RerunSink( + str(output), + filename_template="{shard_id}__w{worker_id}/{segment_id}.rrd", + ) + with pytest.raises(ValueError, match="rendered duplicate output path"): + sink.write_shard_block("shard-a", [unit, unit]) + + def test_write_rerun_without_footer_uses_table_fallback( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, @@ -561,6 +599,29 @@ def test_write_rerun_without_footer_uses_table_fallback( assert row["rerun"].tables["frame"].num_rows == 3 +def test_write_rerun_without_footer_rejects_metadata_only_recording( + tmp_path: Path, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-no-footer-metadata-only" + _tiny_rrd(source) + + row = cast( + Any, + next( + mdr.read_rerun( + str(source), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + sink = RerunSink(str(output), write_footer=False) + + with pytest.raises(ValueError, match="without materialized Rerun table columns"): + sink.write_shard_block("shard-a", [row]) + + def test_write_rerun_table_fallback_separates_static_columns(tmp_path: Path) -> None: source = tmp_path / "sparse.rrd" _sparse_rrd(source) From 9e2f1c01e1673d71a2c0e544536b9bd0dd20b54f Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 03:04:18 +0200 Subject: [PATCH 30/65] Harden Rerun support and cloud metadata --- benchmark/rerun/README.md | 2 + benchmark/rerun/refresh_aws_secrets.py | 3 +- benchmark/rerun/run_cloud_benchmark.py | 26 +- docs/reading-data/rerun.md | 3 + docs/writing-data/rerun.md | 4 +- src/refiner/execution/engine.py | 91 ++++- src/refiner/inference/internal/runtime.py | 28 +- src/refiner/launchers/cloud.py | 6 + src/refiner/pipeline/_rerun_io.py | 81 +++++ src/refiner/pipeline/builtins.py | 76 ++++ src/refiner/pipeline/pipeline.py | 20 +- src/refiner/pipeline/planning.py | 40 +-- src/refiner/pipeline/sinks/base.py | 14 +- src/refiner/pipeline/sinks/rerun.py | 53 +-- src/refiner/pipeline/sources/base.py | 7 +- src/refiner/pipeline/sources/readers/base.py | 2 +- src/refiner/pipeline/sources/readers/rerun.py | 328 ++++++++++++------ src/refiner/platform/client/api.py | 36 +- src/refiner/platform/manifest.py | 18 +- src/refiner/services/discovery.py | 61 +--- tests/launchers/test_cloud_launcher.py | 42 +++ tests/platform/test_client.py | 52 +++ tests/platform/test_client_create_job.py | 29 +- tests/platform/test_cloud_client.py | 70 ++++ tests/platform/test_manifest.py | 35 ++ tests/readers/test_rerun_reader.py | 204 ++++++++++- 26 files changed, 1049 insertions(+), 282 deletions(-) create mode 100644 src/refiner/pipeline/_rerun_io.py create mode 100644 src/refiner/pipeline/builtins.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 4ab23bc3..3fd691ab 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -37,6 +37,8 @@ look better. - Macrodata CLI auth must be configured. - Workspace secrets in the selected environment must include AWS credentials for the source/output S3 bucket. The default environment is `researcher`. + Pass `--aws-profile` explicitly when refreshing those credentials; the helper + intentionally does not fall back to the AWS default profile. If local AWS credentials are valid, refresh the cloud secret environment without printing credential values: diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py index a6254bcd..f8ecd56a 100644 --- a/benchmark/rerun/refresh_aws_secrets.py +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -9,7 +9,6 @@ DEFAULT_S3_CHECK = ( "s3://macrodata-rerun-format-tests/dominique-sample/episode-5__base.rrd" ) -DEFAULT_AWS_PROFILE = "210049840512_Researcher" SECRET_NAMES = ( "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", @@ -25,7 +24,7 @@ def _parse_args() -> argparse.Namespace: "workspace secret environment for Rerun cloud benchmarks." ) ) - parser.add_argument("--aws-profile", default=DEFAULT_AWS_PROFILE) + parser.add_argument("--aws-profile", required=True) parser.add_argument("--secret-env", default="researcher") parser.add_argument( "--region", diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index def679de..93ef658f 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -411,18 +411,24 @@ def _inspect_output(path: str) -> tuple[int | None, int | None, str | None]: fs, fs_path = url_to_fs(path) if not fs.exists(fs_path): return 0, 0, None + root_info = fs.info(fs_path) + if root_info.get("type") != "directory": + return 1, int(root_info.get("size", 0)), None total_size = 0 total_files = 0 - found = fs.find(fs_path, detail=True) - if isinstance(found, dict): - infos = found.values() - else: - infos = (fs.info(child) for child in fs.find(fs_path)) - for info in infos: - if info.get("type") == "directory": - continue - total_files += 1 - total_size += int(info.get("size", 0)) + pending = [fs_path] + while pending: + current = pending.pop() + for info in fs.ls(current, detail=True): + child_type = info.get("type") + child_name = info.get("name") or info.get("Key") + if not isinstance(child_name, str): + continue + if child_type == "directory": + pending.append(child_name) + else: + total_files += 1 + total_size += int(info.get("size", 0)) return total_files, total_size, None except Exception as err: return None, None, str(err) diff --git a/docs/reading-data/rerun.md b/docs/reading-data/rerun.md index deaca638..894df0a4 100644 --- a/docs/reading-data/rerun.md +++ b/docs/reading-data/rerun.md @@ -54,6 +54,9 @@ For raw RRD copy workflows that immediately call `write_rerun`, set `materialize_tables=False` with `output="recording"`. The row still carries the source recording metadata needed by the writer's chunk-copy path, but skips the Arrow timeline/static tables that downstream code will not inspect. +Use this for pure copies; timeline-filtered writes should keep +`materialize_tables=True` so the writer can emit exactly the selected timeline +tables. ## Robotics rows diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index e8874b49..f68fe985 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -47,7 +47,9 @@ Python. For pure copy jobs, use `read_rerun(..., materialize_tables=False)` before `write_rerun(...)` to skip timeline/static table materialization while keeping -the raw source chunks available to the writer. +the raw source chunks available to the writer. If the read applies explicit +timeline filters, keep `materialize_tables=True`; metadata-only rows cannot +project timelines exactly during raw chunk copying. If a `RerunRecording` has no source file, the writer falls back to table emission with `send_dataframe`. Static Rerun component columns are sent as diff --git a/src/refiner/execution/engine.py b/src/refiner/execution/engine.py index 0006fa15..b9bf4f80 100644 --- a/src/refiner/execution/engine.py +++ b/src/refiner/execution/engine.py @@ -9,6 +9,7 @@ from refiner.pipeline.data.block import Block, StreamItem from refiner.pipeline.data.datatype import schema_with_dtypes +from refiner.pipeline.data.shard import SHARD_ID_COLUMN from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.steps import ( CastStep, @@ -31,7 +32,7 @@ from refiner.execution.operators.vectorized import ( apply_vectorized_ops, ) -from refiner.pipeline.data.row import Row +from refiner.pipeline.data.row import DictRow, Row _DEFAULT_VECTORIZED_CHUNK_ROWS = 2048 @@ -301,27 +302,50 @@ def _execute_vector_segment( pending_rows = RowBuffer() current_chunk_rows = max(1, int(vectorized_chunk_rows)) estimated_row_bytes: float | None = None - segment_changes_rows = any( - isinstance(op, (FilterExprStep, FnTableStep)) for op in ops + row_projection_ops, row_remaining_ops = _split_row_projection_ops(ops) + row_projected_schema = ( + _vector_segment_schema(input_schema, row_projection_ops) + if row_projection_ops + else input_schema ) - def _run_block(block: Tabular) -> Tabular: - return_row_indices = block.needs_row_indices and segment_changes_rows + def _run_block(block: Tabular, block_ops: Sequence[VectorizedOp]) -> Tabular: + block_changes_rows = any( + isinstance(op, (FilterExprStep, FnTableStep)) for op in block_ops + ) + return_row_indices = block.needs_row_indices and block_changes_rows if not return_row_indices: table = apply_vectorized_ops( block.table, - ops, + block_ops, on_shard_delta=on_shard_delta, ) return block.with_table(table) table, row_indices = apply_vectorized_ops( block.table, - ops, + block_ops, on_shard_delta=on_shard_delta, return_row_indices=True, ) return block.with_table(table, row_indices=row_indices) + def _rows_to_block(batch: list[Row]) -> tuple[Tabular, Sequence[VectorizedOp]]: + try: + return _tabular_from_rows(batch, schema=input_schema), ops + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as err: + if not row_projection_ops: + raise + projected = [ + _apply_row_projection_ops(row, row_projection_ops) for row in batch + ] + try: + return ( + _tabular_from_rows(projected, schema=row_projected_schema), + row_remaining_ops, + ) + except (pa.ArrowInvalid, pa.ArrowTypeError, TypeError) as fallback_err: + raise err from fallback_err + def _chunk_rows_for_budget() -> int: if ( max_vectorized_block_bytes is None @@ -338,11 +362,7 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: while True: batch = pending_rows.peek(rows_for_try) try: - block = ( - Tabular.from_rows(batch, schema=input_schema) - if not batch - else batch[0].tabular_type.from_rows(batch, schema=input_schema) - ) + block, block_ops = _rows_to_block(batch) except pa.ArrowMemoryError: if rows_for_try <= 1: raise @@ -368,7 +388,7 @@ def _run_pending_chunk(target_rows: int) -> Iterator[Tabular]: continue try: - out = _run_block(block) + out = _run_block(block, block_ops) except pa.ArrowMemoryError: if rows_for_try <= 1: raise @@ -420,7 +440,7 @@ def _yield_tabular_chunks(block: Tabular) -> Iterator[Tabular]: continue try: - out = _run_block(chunk) + out = _run_block(chunk, ops) except pa.ArrowMemoryError: if chunk_rows <= 1: raise @@ -457,6 +477,49 @@ def _yield_tabular_chunks(block: Tabular) -> Iterator[Tabular]: yield from _drain_rows(force=True) +def _tabular_from_rows( + rows: list[Row], + *, + schema: pa.Schema | None, +) -> Tabular: + return ( + Tabular.from_rows(rows, schema=schema) + if not rows + else rows[0].tabular_type.from_rows(rows, schema=schema) + ) + + +def _split_row_projection_ops( + ops: Sequence[VectorizedOp], +) -> tuple[tuple[SelectStep | DropStep, ...], Sequence[VectorizedOp]]: + projection: list[SelectStep | DropStep] = [] + for index, op in enumerate(ops): + if not isinstance(op, (SelectStep, DropStep)): + return tuple(projection), ops[index:] + projection.append(op) + return tuple(projection), () + + +def _apply_row_projection_ops( + row: Row, + ops: Sequence[SelectStep | DropStep], +) -> Row: + out = row + for op in ops: + if isinstance(op, SelectStep): + out = DictRow( + { + column: out[column] + for column in op.columns + if column != SHARD_ID_COLUMN + }, + shard_id=out.shard_id, + ) + continue + out = out.drop(*(column for column in op.columns if column != SHARD_ID_COLUMN)) + return out + + def _chunk_output_rows(rows: Iterable[Row], block_rows: int) -> Iterator[list[Row]]: pending: list[Row] = [] for row in rows: diff --git a/src/refiner/inference/internal/runtime.py b/src/refiner/inference/internal/runtime.py index 9f955603..1f41063c 100644 --- a/src/refiner/inference/internal/runtime.py +++ b/src/refiner/inference/internal/runtime.py @@ -18,14 +18,13 @@ _OpenAIResponsesClient, ) from refiner.inference.types import InferenceProvider +from refiner.pipeline.builtins import REFINER_BUILTIN_CALL_ATTR, builtin_call_spec from refiner.pipeline.data.row import Row from refiner.pipeline.steps import MapResult from refiner.services import VLLMRuntimeServiceBinding from refiner.worker.context import get_active_service_manager from refiner.worker.metrics.api import register_gauge -_REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__" - RequestFn: TypeAlias = Callable[[Mapping[str, Any]], Awaitable[Any]] MapFn: TypeAlias = Callable[[Row, RequestFn], Awaitable[MapResult] | MapResult] ClientCall: TypeAlias = Callable[[Any, Mapping[str, Any]], Awaitable[Any]] @@ -168,27 +167,30 @@ async def _close() -> None: } if defaults_key is not None: args[defaults_key] = dict(defaults or {}) - builtin = getattr(fn, _REFINER_BUILTIN_CALL_ATTR, None) + builtin = builtin_call_spec(fn) builtin_name = name builtin_args = args - if isinstance(builtin, dict): - candidate_name = builtin.get("name") - candidate_args = builtin.get("args") - if isinstance(candidate_name, str) and candidate_name: - builtin_name = candidate_name - if isinstance(candidate_args, dict): - builtin_args = candidate_args + builtin_services = [] + builtin_refiner_extras: tuple[str, ...] = () + if builtin is not None: + builtin_name = builtin.name + builtin_args = builtin.args + builtin_services.extend(builtin.services) + builtin_refiner_extras = builtin.refiner_extras + if service is not None: + builtin_services.append(service.to_spec()) setattr( _wrapped, - _REFINER_BUILTIN_CALL_ATTR, + REFINER_BUILTIN_CALL_ATTR, { "name": builtin_name, "args": builtin_args, - "services": [] if service is None else [service.to_spec()], + "services": builtin_services, + "refiner_extras": builtin_refiner_extras, }, ) setattr(_wrapped, "aclose", _close) return _wrapped -__all__ = ["_REFINER_BUILTIN_CALL_ATTR", "inference_map"] +__all__ = ["REFINER_BUILTIN_CALL_ATTR", "inference_map"] diff --git a/src/refiner/launchers/cloud.py b/src/refiner/launchers/cloud.py index a1e17b99..38d6d1b9 100644 --- a/src/refiner/launchers/cloud.py +++ b/src/refiner/launchers/cloud.py @@ -105,6 +105,7 @@ class CloudLauncher(BaseLauncher): local environment in the cloud runtime. dependencies: Additional packages to install in the cloud runtime. Entries are requirement strings. + extra_dependencies: Compatibility alias for ``dependencies``. refiner_extras: Additional macrodata-refiner extras to install in the cloud runtime. Built-in blocks automatically declare the extras they require; pass this for extras used outside those blocks. @@ -123,6 +124,7 @@ def __init__( gpu: GPU | None = None, sync_local_dependencies: bool = False, dependencies: Sequence[str] | None = None, + extra_dependencies: Sequence[str] | None = None, refiner_extras: Sequence[str] | None = None, secrets: SecretInput | None = None, env: dict[str, object | None] | None = None, @@ -141,6 +143,10 @@ def __init__( raise ValueError("unsafe_continue requires continue_from_job") if mem_mb_per_worker is not None and mem_mb_per_worker <= 0: raise ValueError("mem_mb_per_worker must be > 0") + if dependencies is not None and extra_dependencies is not None: + raise ValueError("Pass only one of dependencies or extra_dependencies") + if dependencies is None: + dependencies = extra_dependencies self.cpus_per_worker = cpus_per_worker self.mem_mb_per_worker = mem_mb_per_worker self.sync_local_dependencies = sync_local_dependencies diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py new file mode 100644 index 00000000..79e7ea7c --- /dev/null +++ b/src/refiner/pipeline/_rerun_io.py @@ -0,0 +1,81 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +import os +import tempfile +from pathlib import Path +from typing import cast + +from refiner.io import DataFile +from refiner.pipeline.data.tabular import Tabular + + +class LocalRrd: + def __init__(self, source: DataFile) -> None: + self.source = source + self.tmpdir: tempfile.TemporaryDirectory[str] | None = None + self.path: Path | None = None + + def open(self) -> Path: + if self.path is not None: + return self.path + if self.source.is_local: + self.path = Path(self.source.abs_path()) + return self.path + self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") + name = os.path.basename(self.source.path) or "recording.rrd" + self.path = Path(self.tmpdir.name) / name + self.source.copy(str(self.path)) + return self.path + + def close(self) -> None: + if self.tmpdir is not None: + self.tmpdir.cleanup() + self.tmpdir = None + self.path = None + + def __enter__(self) -> Path: + return self.open() + + def __exit__(self, *args: object) -> None: + self.close() + + def __del__(self) -> None: + try: + self.close() + except Exception: + pass + + def __getstate__(self) -> dict[str, object]: + return { + "source": self.source, + "path": str(self.path) if self.path is not None else None, + } + + def __setstate__(self, state: dict[str, object]) -> None: + self.source = cast(DataFile, state["source"]) + self.tmpdir = None + path = state.get("path") + self.path = Path(path) if isinstance(path, str) else None + + +@dataclass(frozen=True, slots=True) +class RerunRecording: + """Columnar Rerun recording data loaded from one RRD segment.""" + + segment_id: str + source_path: str + tables: Mapping[str, Tabular] + static: Tabular | None = None + source_file: DataFile | None = None + local_source: LocalRrd | None = None + application_id: str | None = None + recording_id: str | None = None + contents: tuple[str, ...] | None = None + timelines: tuple[str, ...] | None = None + include_static: bool = True + use_source_chunks: bool = True + + +__all__ = ["LocalRrd", "RerunRecording"] diff --git a/src/refiner/pipeline/builtins.py b/src/refiner/pipeline/builtins.py new file mode 100644 index 00000000..304682ab --- /dev/null +++ b/src/refiner/pipeline/builtins.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +from refiner.pipeline.steps import VectorizedSegmentStep +from refiner.services.base import RuntimeServiceSpec + +if TYPE_CHECKING: + from refiner.pipeline import RefinerPipeline + +REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__" + + +@dataclass(frozen=True, slots=True) +class BuiltinCallSpec: + name: str + args: dict[str, Any] + services: tuple[RuntimeServiceSpec, ...] = () + refiner_extras: tuple[str, ...] = () + + +def builtin_call_spec(fn: Any) -> BuiltinCallSpec | None: + spec = getattr(fn, REFINER_BUILTIN_CALL_ATTR, None) + if not isinstance(spec, dict): + return None + name = spec.get("name") + if not isinstance(name, str) or not name: + return None + args = spec.get("args") + if not isinstance(args, dict): + return None + services = spec.get("services", ()) + if not isinstance(services, (list, tuple)): + return None + parsed_services: list[RuntimeServiceSpec] = [] + for service in services: + if not isinstance(service, RuntimeServiceSpec): + return None + parsed_services.append(service) + refiner_extras = spec.get("refiner_extras", ()) + if not isinstance(refiner_extras, tuple) or not all( + isinstance(extra, str) for extra in refiner_extras + ): + return None + return BuiltinCallSpec( + name=name, + args=args, + services=tuple(parsed_services), + refiner_extras=refiner_extras, + ) + + +def iter_pipeline_builtin_specs( + pipeline: "RefinerPipeline", +) -> Iterator[BuiltinCallSpec]: + seen: set[int] = set() + for step in pipeline.pipeline_steps: + candidates = step.ops if isinstance(step, VectorizedSegmentStep) else (step,) + for candidate in candidates: + for attr in ("fn", "predicate"): + fn = getattr(candidate, attr, None) + if fn is None or id(fn) in seen: + continue + seen.add(id(fn)) + if spec := builtin_call_spec(fn): + yield spec + + +__all__ = [ + "BuiltinCallSpec", + "REFINER_BUILTIN_CALL_ATTR", + "builtin_call_spec", + "iter_pipeline_builtin_specs", +] diff --git a/src/refiner/pipeline/pipeline.py b/src/refiner/pipeline/pipeline.py index 2764f7c9..4cc0ae51 100644 --- a/src/refiner/pipeline/pipeline.py +++ b/src/refiner/pipeline/pipeline.py @@ -49,7 +49,12 @@ from refiner.pipeline.sources.readers.hdf5 import MissingPolicy from refiner.pipeline.sources.readers.lerobot import LeRobotEpisodeReader from refiner.pipeline.sources.readers.mcap import SyncMethod -from refiner.pipeline.sources.readers.rerun import RerunOutputMode +from refiner.pipeline.sources.readers.rerun import ( + DEFAULT_RERUN_ACTION_PREFIX, + DEFAULT_RERUN_CAMERA_PREFIX, + DEFAULT_RERUN_STATE_PREFIX, + RerunOutputMode, +) from refiner.pipeline.sources.items import ItemsSource from refiner.pipeline.sources.task import TaskSource, TaskStep from refiner.pipeline.data import datatype @@ -736,6 +741,7 @@ def launch_cloud( gpu: GPU | None = None, sync_local_dependencies: bool = False, dependencies: Sequence[str] | None = None, + extra_dependencies: Sequence[str] | None = None, refiner_extras: Sequence[str] | None = None, secrets: SecretInput | None = None, env: Mapping[str, object | None] | None = None, @@ -755,6 +761,7 @@ def launch_cloud( dependencies: Additional packages to install in the cloud runtime. Entries are requirement strings such as `"torch"` or `"ego-vision[models]==0.1.2"`. + extra_dependencies: Compatibility alias for ``dependencies``. refiner_extras: Additional macrodata-refiner extras to install in the cloud runtime. Built-in blocks automatically declare the extras they require; pass this for extras used outside those @@ -772,6 +779,11 @@ def launch_cloud( """ from refiner.launchers.cloud import CloudLauncher + if dependencies is not None and extra_dependencies is not None: + raise ValueError("Pass only one of dependencies or extra_dependencies") + if dependencies is None: + dependencies = extra_dependencies + launcher = CloudLauncher( pipeline=self, name=name, @@ -1260,9 +1272,9 @@ def read_rerun( materialize_tables: bool = True, include_recording: bool | None = None, fill_latest_at: bool = False, - action_prefix: str = "/action", - state_prefix: str = "/observation/state", - camera_prefix: str = "/cam", + action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, + state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, + camera_prefix: str = DEFAULT_RERUN_CAMERA_PREFIX, actions: PathSelection | None = None, states: PathSelection | None = None, videos: PathSelection | None = None, diff --git a/src/refiner/pipeline/planning.py b/src/refiner/pipeline/planning.py index 2e470779..c789deed 100644 --- a/src/refiner/pipeline/planning.py +++ b/src/refiner/pipeline/planning.py @@ -22,10 +22,10 @@ VectorizedSegmentStep, WithColumnsStep, ) +from refiner.pipeline.builtins import REFINER_BUILTIN_CALL_ATTR, builtin_call_spec from refiner.pipeline.data.datatype import dtype_to_plan from refiner.pipeline.resources import GPU from refiner.platform.manifest import _redact_captured_text -from refiner.services import RuntimeServiceSpec from refiner.services.discovery import ( collect_pipeline_services, runtime_service_specs_to_dicts, @@ -35,9 +35,6 @@ from refiner.pipeline import RefinerPipeline -_REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__" - - @dataclass(frozen=True, slots=True) class StageComputeRequirements: num_workers: int @@ -66,9 +63,9 @@ class PlannedStage: def _explicit_callable_name(fn: Any) -> str | None: - builtin_description = _builtin_description(fn) - if builtin_description is not None: - return builtin_description["name"] + spec = builtin_call_spec(fn) + if spec is not None: + return spec.name name = getattr(fn, "__name__", None) if not isinstance(name, str): return None @@ -84,13 +81,13 @@ def _callable_step_args( extra_args: dict[str, Any] | None = None, builtin_extra_args: dict[str, Any] | None = None, ) -> dict[str, Any]: - builtin_description = _builtin_description(fn) - if builtin_description is None: + spec = builtin_call_spec(fn) + if spec is None: args: dict[str, Any] = {"fn": fn} if extra_args: args.update(extra_args) else: - args = dict(builtin_description["args"]) + args = dict(spec.args) if builtin_extra_args: args.update(builtin_extra_args) return args @@ -310,34 +307,13 @@ def _callable_source(fn: Any) -> str: return repr(fn) -def _builtin_description(fn: Any) -> dict[str, Any] | None: - spec = getattr(fn, _REFINER_BUILTIN_CALL_ATTR, None) - if not isinstance(spec, dict): - return None - name = spec.get("name") - if not isinstance(name, str) or not name: - return None - args = spec.get("args") - if not isinstance(args, dict): - return None - services = spec.get("services", ()) - if not isinstance(services, (list, tuple)): - return None - parsed_services: list[RuntimeServiceSpec] = [] - for service in services: - if not isinstance(service, RuntimeServiceSpec): - return None - parsed_services.append(service) - return {"name": name, "args": args, "services": tuple(parsed_services)} - - def describe_builtin( name: str, *, refiner_extras: tuple[str, ...] = (), **args: Any ) -> Any: def _decorate(fn: Any) -> Any: setattr( fn, - _REFINER_BUILTIN_CALL_ATTR, + REFINER_BUILTIN_CALL_ATTR, { "name": name, "args": args, diff --git a/src/refiner/pipeline/sinks/base.py b/src/refiner/pipeline/sinks/base.py index 1a1776be..fce688c6 100644 --- a/src/refiner/pipeline/sinks/base.py +++ b/src/refiner/pipeline/sinks/base.py @@ -72,10 +72,16 @@ def _declared_refiner_extras(self) -> tuple[str, ...]: return () def _io_refiner_extras(self) -> tuple[str, ...]: - """Storage extras required by this sink's output, if it has one.""" - if not hasattr(self, "output"): - return () - return cast(Any, self).output.required_refiner_extras() + """Storage extras required by this sink's IO handles. + + Keep a structural fallback for existing custom sinks that expose an + ``output`` object with ``required_refiner_extras()``. + """ + output = getattr(self, "output", None) + required_refiner_extras = getattr(output, "required_refiner_extras", None) + if callable(required_refiner_extras): + return tuple(cast(Any, required_refiner_extras)()) + return () def build_reducer(self) -> "BaseSink | None": """Return an optional 1-worker reducer sink for launched execution. diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index 4a8e4344..036191ab 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -2,6 +2,7 @@ import os import tempfile +import warnings from pathlib import Path from string import Formatter from typing import Any @@ -10,11 +11,11 @@ from refiner.io.datafile import DataFile from refiner.io.datafolder import DataFolder, DataFolderLike +from refiner.pipeline._rerun_io import LocalRrd, RerunRecording from refiner.pipeline.data.block import Block from refiner.pipeline.data.row import Row from refiner.pipeline.sinks.base import BaseSink from refiner.pipeline.sinks.reducer.file import FileCleanupReducerSink -from refiner.pipeline.sources.readers.rerun import RerunRecording, _local_rrd from refiner.utils import check_required_dependencies from refiner.worker.context import get_active_worker_token from refiner.worker.metrics.api import log_throughput @@ -33,13 +34,14 @@ def __init__( app_id: str = "refiner", write_footer: bool = True, ) -> None: - _validate_filename_template(filename_template) + template_fields = _validate_filename_template(filename_template) self.output = DataFolder.resolve(output) self.filename_template = filename_template + self._uses_segment_id = "segment_id" in template_fields self.app_id = app_id self.write_footer = write_footer self._row_indices: dict[str, int] = {} - self._written_relpaths: set[str] = set() + self._written_relpaths: dict[str, set[str]] = {} def _declared_refiner_extras(self) -> tuple[str, ...]: return ("rerun",) @@ -56,16 +58,19 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: worker_id=get_active_worker_token(), row_index=row_index, segment_id=recording.segment_id, + uses_segment_id=self._uses_segment_id, ) - if relpath in self._written_relpaths: + written_relpaths = self._written_relpaths.setdefault(shard_id, set()) + if relpath in written_relpaths: raise ValueError( "write_rerun filename_template rendered duplicate output path " f"{relpath!r}; include {{row_index}} or another unique row field" ) self._write_recording(recording, relpath) - self._written_relpaths.add(relpath) - log_throughput("files_written", 1, shard_id=shard_id, unit="files") + written_relpaths.add(relpath) count += 1 + if count: + log_throughput("files_written", count, shard_id=shard_id, unit="files") return count def _write_recording(self, recording: RerunRecording, relpath: str) -> None: @@ -104,6 +109,7 @@ def write_local(path: Path) -> None: def on_shard_complete(self, shard_id: str) -> None: self._row_indices.pop(shard_id, None) + self._written_relpaths.pop(shard_id, None) def describe(self) -> tuple[str, str, dict[str, object]]: return ( @@ -138,10 +144,20 @@ def _write_source_chunks( *, application_id: str, ) -> None: + local_source = recording.local_source + local_source_path = local_source.path if local_source is not None else None + if local_source_path is not None and local_source_path.exists(): + _write_source_chunks_from_path( + recording, + path, + local_path=local_source_path, + application_id=application_id, + ) + return source = recording.source_file if source is None: raise ValueError("Rerun source chunk write requires source_file") - with _local_rrd(source) as local_path: + with LocalRrd(source) as local_path: _write_source_chunks_from_path( recording, path, @@ -159,7 +175,12 @@ def _write_source_chunks_from_path( ) -> None: import rerun as rr - reader = rr.experimental.RrdReader(local_path) + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="RRD file has no footer/manifest:.*", + ) + reader = rr.experimental.RrdReader(local_path) store = _matching_store(reader, recording) stream = reader.stream(store=store) if recording.contents is not None: @@ -294,7 +315,7 @@ def _is_static_column(field: pa.Field) -> bool: return (field.metadata or {}).get(b"rerun:is_static") == b"true" -def _validate_filename_template(filename_template: str) -> None: +def _validate_filename_template(filename_template: str) -> set[str]: fields: set[str] = set() for _literal_text, field_name, format_spec, conversion in Formatter().parse( filename_template @@ -329,6 +350,7 @@ def _validate_filename_template(filename_template: str) -> None: ), "filename_template", ) + return fields def _render_relpath( @@ -338,6 +360,7 @@ def _render_relpath( worker_id: str, row_index: int, segment_id: str, + uses_segment_id: bool, ) -> str: field_values: dict[str, object] = { "shard_id": shard_id, @@ -345,7 +368,7 @@ def _render_relpath( "row_index": row_index, "segment_id": segment_id, } - if "segment_id" in _template_fields(filename_template): + if uses_segment_id: field_values["segment_id"] = _normalize_path_segment(segment_id, "segment_id") return _normalize_relpath( filename_template.format(**field_values), @@ -353,16 +376,6 @@ def _render_relpath( ) -def _template_fields(filename_template: str) -> set[str]: - return { - field_name - for _literal_text, field_name, _format_spec, _conversion in Formatter().parse( - filename_template - ) - if field_name is not None - } - - def _normalize_path_segment(value: str, label: str) -> str: if not value or value in {".", ".."} or "/" in value or "\\" in value: raise ValueError(f"{label} must be a single relative path segment") diff --git a/src/refiner/pipeline/sources/base.py b/src/refiner/pipeline/sources/base.py index 75e85b52..cff34ee3 100644 --- a/src/refiner/pipeline/sources/base.py +++ b/src/refiner/pipeline/sources/base.py @@ -12,7 +12,7 @@ from refiner.worker.metrics.api import log_throughput _INTERNAL_SHARD_ID_KEY = "__shard_id" -SourceUnit: TypeAlias = Row | Tabular +SourceUnit: TypeAlias = Row | list[Row] | Tabular class BaseSource(ABC): @@ -73,6 +73,8 @@ def _io_refiner_extras(self) -> tuple[str, ...]: def _unit_num_rows(unit: SourceUnit) -> int: if isinstance(unit, Row): return 1 + if isinstance(unit, list): + return len(unit) if isinstance(unit, Tabular): return int(unit.num_rows) raise TypeError(f"Unsupported source unit type: {type(unit)!r}") @@ -82,6 +84,9 @@ def _with_shard_id(unit: SourceUnit, shard_id: str) -> SourceUnit: if isinstance(unit, Row): return unit.update(**{_INTERNAL_SHARD_ID_KEY: shard_id}) + if isinstance(unit, list): + return [row.update(**{_INTERNAL_SHARD_ID_KEY: shard_id}) for row in unit] + if isinstance(unit, Tabular): table = unit.table if table.num_rows == 0: diff --git a/src/refiner/pipeline/sources/readers/base.py b/src/refiner/pipeline/sources/readers/base.py index 5464680f..69bfa722 100644 --- a/src/refiner/pipeline/sources/readers/base.py +++ b/src/refiner/pipeline/sources/readers/base.py @@ -317,7 +317,7 @@ def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: Contract: - Must accept shards returned by `list_shards()`. - Should be safe to call sequentially (single-worker, no concurrent calls). - - Units can be `Row` or `Tabular`. + - Units can be `Row`, row blocks, or `Tabular`. """ raise NotImplementedError diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 0998b0f0..00d8d33f 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -1,12 +1,9 @@ from __future__ import annotations -import os -import tempfile -from contextlib import ExitStack -from collections.abc import Iterator, Mapping, Sequence -from dataclasses import dataclass +from collections.abc import Iterable, Iterator, Mapping, Sequence from pathlib import Path from typing import Any, Literal, cast +import warnings import numpy as np import pyarrow as pa @@ -15,7 +12,8 @@ from refiner.io import DataFile from refiner.io.fileset import DataFileSetLike -from refiner.pipeline.data.row import DictRow +from refiner.pipeline._rerun_io import LocalRrd, RerunRecording +from refiner.pipeline.data.row import DictRow, Row from refiner.pipeline.data.shard import FilePartsDescriptor, Shard from refiner.pipeline.data.tabular import Tabular from refiner.pipeline.sources.base import SourceUnit @@ -27,6 +25,7 @@ ) from refiner.utils import check_required_dependencies from refiner.video import VideoFrameSequence +from refiner.worker.context import logger RerunOutputMode = Literal["recording", "robotics"] @@ -37,23 +36,50 @@ {"episode_id", "rerun", "frames", "fps", "robot_type"} ) _RECORDING_ROW_COLUMNS = frozenset({"episode_id", "rerun"}) +DEFAULT_RERUN_ACTION_PREFIX = "/action" +DEFAULT_RERUN_STATE_PREFIX = "/observation/state" +DEFAULT_RERUN_CAMERA_PREFIX = "/cam" +# Amortize Rerun server startup for small files without staging an unbounded shard. +_MAX_STAGED_RRD_BATCH_BYTES = 512 * 1024 * 1024 +_MAX_STAGED_RRD_BATCH_FILES = 16 -@dataclass(frozen=True, slots=True) -class RerunRecording: - """Columnar Rerun recording data loaded from one RRD segment.""" - - segment_id: str - source_path: str - tables: Mapping[str, Tabular] - static: Tabular | None = None - source_file: DataFile | None = None - application_id: str | None = None - recording_id: str | None = None - contents: tuple[str, ...] | None = None - timelines: tuple[str, ...] | None = None - include_static: bool = True - use_source_chunks: bool = True +def _reject_recording_robotics_options( + *, + primary_timeline: str | None, + action_prefix: str, + state_prefix: str, + camera_prefix: str, + actions: PathSelection | None, + states: PathSelection | None, + videos: PathSelection | None, + fps: float | None, + robot_type: str | None, +) -> None: + invalid = [] + if primary_timeline is not None: + invalid.append("primary_timeline") + if action_prefix != DEFAULT_RERUN_ACTION_PREFIX: + invalid.append("action_prefix") + if state_prefix != DEFAULT_RERUN_STATE_PREFIX: + invalid.append("state_prefix") + if camera_prefix != DEFAULT_RERUN_CAMERA_PREFIX: + invalid.append("camera_prefix") + if actions is not None: + invalid.append("actions") + if states is not None: + invalid.append("states") + if videos is not None: + invalid.append("videos") + if fps is not None: + invalid.append("fps") + if robot_type is not None: + invalid.append("robot_type") + if invalid: + raise ValueError( + "Rerun recording output does not use robotics options: " + + ", ".join(invalid) + ) class RerunReader(BaseReader): @@ -79,9 +105,9 @@ def __init__( materialize_tables: bool = True, include_recording: bool | None = None, fill_latest_at: bool = False, - action_prefix: str = "/action", - state_prefix: str = "/observation/state", - camera_prefix: str = "/cam", + action_prefix: str = DEFAULT_RERUN_ACTION_PREFIX, + state_prefix: str = DEFAULT_RERUN_STATE_PREFIX, + camera_prefix: str = DEFAULT_RERUN_CAMERA_PREFIX, actions: PathSelection | None = None, states: PathSelection | None = None, videos: PathSelection | None = None, @@ -98,6 +124,18 @@ def __init__( raise ValueError( "include_recording=False is only supported for robotics output" ) + if output == "recording": + _reject_recording_robotics_options( + primary_timeline=primary_timeline, + action_prefix=action_prefix, + state_prefix=state_prefix, + camera_prefix=camera_prefix, + actions=actions, + states=states, + videos=videos, + fps=fps, + robot_type=robot_type, + ) if fps is not None: fps = float(fps) if not np.isfinite(fps) or fps <= 0: @@ -119,6 +157,7 @@ def __init__( self.primary_timeline = primary_timeline self.include_static = include_static self.materialize_tables = materialize_tables + self.use_source_chunks = self.timelines is None self.include_recording = ( output == "recording" if include_recording is None else include_recording ) @@ -171,36 +210,80 @@ def describe(self) -> dict[str, Any]: "output": self.output, "contents": self.contents, "timelines": self.timelines, - "primary_timeline": self.primary_timeline, "include_static": self.include_static, "materialize_tables": self.materialize_tables, "include_recording": self.include_recording, "fill_latest_at": self.fill_latest_at, - "action_prefix": self.action_prefix, - "state_prefix": self.state_prefix, - "camera_prefix": self.camera_prefix, - "actions": dict(self.actions) if self.actions_explicit else None, - "states": dict(self.states) if self.states_explicit else None, - "videos": dict(self.videos) if self.videos_explicit else None, - "fps": self.fps, - "robot_type": self.robot_type, } ) + if self.output == "robotics": + description.update( + { + "primary_timeline": self.primary_timeline, + "action_prefix": self.action_prefix, + "state_prefix": self.state_prefix, + "camera_prefix": self.camera_prefix, + "actions": dict(self.actions) if self.actions_explicit else None, + "states": dict(self.states) if self.states_explicit else None, + "videos": dict(self.videos) if self.videos_explicit else None, + "fps": self.fps, + "robot_type": self.robot_type, + } + ) return description def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: descriptor = shard.descriptor assert isinstance(descriptor, FilePartsDescriptor) - with ExitStack() as stack: - local_files = [] - for part in descriptor.parts: - source = self.fileset.resolve_file(part.source_index, part.path) - local_files.append((source, stack.enter_context(_local_rrd(source)))) + batch: list[tuple[DataFile, Path, LocalRrd]] = [] + batch_bytes = 0 + for part in descriptor.parts: + source = self.fileset.resolve_file(part.source_index, part.path) + part_size = max(0, self.fileset.size(part.source_index, part.path)) + if batch and ( + len(batch) >= _MAX_STAGED_RRD_BATCH_FILES + or batch_bytes + part_size > _MAX_STAGED_RRD_BATCH_BYTES + ): + yield from self._read_staged_batch(batch) + batch = [] + batch_bytes = 0 + local_source = LocalRrd(source) + try: + batch.append((source, local_source.open(), local_source)) + except BaseException: + local_source.close() + raise + batch_bytes += part_size + if batch: + yield from self._read_staged_batch(batch) + + def _read_staged_batch( + self, + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], + ) -> Iterator[SourceUnit]: + if self._retain_batch_local_sources(): + try: + units = list(self._read_files(local_files)) + except BaseException: + _close_local_sources(local_files) + raise + if not units: + _close_local_sources(local_files) + return + try: + yield cast(list[Row], units) + finally: + _close_local_sources(local_files) + return + + try: yield from self._read_files(local_files) + finally: + _close_local_sources(local_files) def _read_files( self, - local_files: Sequence[tuple[DataFile, Path]], + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], ) -> Iterator[SourceUnit]: if self.output == "recording" and not self.materialize_tables: check_required_dependencies( @@ -208,13 +291,17 @@ def _read_files( [("rerun", "rerun-sdk")], dist="rerun", ) - server_fallback: list[tuple[DataFile, Path]] = [] - for source, local_path in local_files: - rows = self._read_metadata_only_recording_rows(source, local_path) + server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] + for source, local_path, local_source in local_files: + rows = self._read_metadata_only_recording_rows( + source, + local_path, + local_source, + ) if rows: yield from rows else: - server_fallback.append((source, local_path)) + server_fallback.append((source, local_path, local_source)) if server_fallback: yield from self._read_files_with_server(server_fallback) return @@ -223,7 +310,7 @@ def _read_files( def _read_files_with_server( self, - local_files: Sequence[tuple[DataFile, Path]], + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], ) -> Iterator[SourceUnit]: check_required_dependencies( "read_rerun", @@ -234,20 +321,26 @@ def _read_files_with_server( datasets = { f"recording_{index}": (str(local_path),) - for index, (_source, local_path) in enumerate(local_files) + for index, (_source, local_path, _local_rrd) in enumerate(local_files) } with rr.server.Server(datasets=cast(Any, datasets)) as server: client = server.client() - for dataset_name, (source, local_path) in zip( + for dataset_name, (source, local_path, local_source) in zip( datasets, local_files, strict=True ): dataset = client.get_dataset(dataset_name) - yield from self._read_dataset(source, local_path, dataset) + yield from self._read_dataset( + source, + local_path, + local_source, + dataset, + ) def _read_dataset( self, source: DataFile, local_path: Path, + local_source: LocalRrd, dataset: Any, ) -> Iterator[SourceUnit]: store_entries = ( @@ -270,6 +363,7 @@ def _read_dataset( segment_id=segment_id, source_path=source.abs_path(), source_file=source, + local_source=local_source, application_id=application_id, recording_id=recording_id, timelines=timelines, @@ -298,7 +392,7 @@ def _read_dataset( yield self._recording_row( segment_id=segment_id, source=source, - local_path=local_path, + local_source=local_source, tables=tables, static=Tabular(static) if static is not None else None, application_id=application_id, @@ -309,6 +403,7 @@ def _read_metadata_only_recording_rows( self, source: DataFile, local_path: Path, + local_source: LocalRrd, ) -> list[DictRow]: rows = [] for store in _recording_entries(local_path): @@ -317,7 +412,7 @@ def _read_metadata_only_recording_rows( self._recording_row( segment_id=recording_id, source=source, - local_path=local_path, + local_source=local_source, tables={}, static=None, application_id=store.application_id, @@ -331,7 +426,7 @@ def _recording_row( *, segment_id: str, source: DataFile, - local_path: Path, + local_source: LocalRrd, tables: Mapping[str, Tabular], static: Tabular | None, application_id: str | None, @@ -345,19 +440,21 @@ def _recording_row( tables=tables, static=static, source_file=source, + local_source=( + local_source if self._retain_batch_local_sources() else None + ), application_id=application_id, recording_id=recording_id, contents=self.contents, timelines=self.timelines, include_static=self.include_static, + use_source_chunks=self.use_source_chunks, ), } self._with_file_path(data, source) return DictRow(data) def _timelines(self, schema: Any) -> tuple[str, ...]: - if self.timelines is not None: - return self.timelines return tuple(str(index.name) for index in schema.index_columns()) def _primary_timeline(self, timelines: Sequence[str]) -> str: @@ -400,6 +497,7 @@ def _robotics_row( segment_id: str, source_path: str, source_file: DataFile, + local_source: LocalRrd, application_id: str | None, recording_id: str, timelines: Sequence[str], @@ -429,11 +527,15 @@ def _robotics_row( tables={timeline: Tabular(table)}, static=Tabular(static) if static is not None else None, source_file=source_file, + local_source=( + local_source if self._retain_batch_local_sources() else None + ), application_id=application_id, recording_id=recording_id, contents=tuple(contents), timelines=(timeline,), include_static=self.include_static, + use_source_chunks=False, ) if self.fps is not None: row["fps"] = self.fps @@ -495,6 +597,13 @@ def _robotics_row( row[self.file_path_column] = source_path return DictRow(row) + def _retain_batch_local_sources(self) -> bool: + return ( + self.output == "recording" + and not self.materialize_tables + and self.use_source_chunks + ) + def _contents(contents: str | Sequence[str] | None) -> tuple[str, ...] | None: if contents is None: @@ -541,35 +650,31 @@ def _collect_table(df: Any) -> pa.Table: return table +def _close_local_sources( + local_files: Iterable[tuple[DataFile, Path, LocalRrd]], +) -> None: + for _source, _local_path, local_source in local_files: + local_source.close() + + def _recording_entries(local_path: Path) -> list[Any]: import rerun as rr try: - return list(rr.experimental.RrdReader(local_path).recordings()) - except Exception: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="RRD file has no footer/manifest:.*", + ) + return list(rr.experimental.RrdReader(local_path).recordings()) + except Exception as err: + logger.warning( + "Rerun recording metadata unavailable; falling back to server scan: {}", + type(err).__name__, + ) return [] -class _local_rrd: - def __init__(self, source: DataFile) -> None: - self.source = source - self.tmpdir: tempfile.TemporaryDirectory[str] | None = None - self.path: Path | None = None - - def __enter__(self) -> Path: - if self.source.is_local: - return Path(self.source.abs_path()) - self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") - name = os.path.basename(self.source.path) or "recording.rrd" - self.path = Path(self.tmpdir.name) / name - self.source.copy(str(self.path)) - return self.path - - def __exit__(self, *args: object) -> None: - if self.tmpdir is not None: - self.tmpdir.cleanup() - - def _metadata_text(metadata: Mapping[bytes, bytes], key: bytes) -> str | None: value = metadata.get(key) return value.decode("utf-8") if value is not None else None @@ -622,10 +727,20 @@ def _matches_entity_prefix(entity_path: str, prefix: str) -> bool: def _camera_columns(columns: Mapping[str, str], prefix: str) -> dict[str, str]: out: dict[str, str] = {} + output_paths: dict[str, str] = {} for entity_path, column in columns.items(): if not _matches_entity_prefix(entity_path, prefix): continue - out[entity_path.strip("/").replace("/", ".")] = column + name = entity_path.strip("/").replace("/", ".") + existing = output_paths.get(name) + if existing is not None: + raise ValueError( + "Rerun camera paths derive the same output video name " + f"{name!r}: {existing!r} and {entity_path!r}; pass explicit " + "videos={...} to choose unique names" + ) + output_paths[name] = entity_path + out[name] = column return out @@ -729,41 +844,56 @@ def _iter_encoded_images(values: pa.Array) -> Iterator[np.ndarray]: from PIL import Image - for index in range(len(values)): - data = _encoded_image_bytes(values, index) - if data is None: - continue - with Image.open(BytesIO(data)) as image: - yield np.asarray(image.convert("RGB"), dtype=np.uint8) - - -def _encoded_image_bytes(values: pa.Array, index: int) -> bytes | None: if not pa.types.is_list(values.type) and not pa.types.is_large_list(values.type): raise TypeError( f"Expected a Rerun encoded image list column, got {values.type}" ) - if not values[index].is_valid: - return None outer_offsets = np.asarray(values.offsets) - outer_start = int(outer_offsets[index]) - outer_end = int(outer_offsets[index + 1]) - if outer_end <= outer_start: - return None + valid = np.asarray(_is_valid(values), dtype=bool) inner = values.values - if not pa.types.is_list(inner.type) and not pa.types.is_large_list(inner.type): + if pa.types.is_list(inner.type) or pa.types.is_large_list(inner.type): + inner_offsets = np.asarray(inner.offsets) + inner_values = inner.values + for index in range(len(values)): + if not valid[index]: + continue + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + continue + byte_start = int(inner_offsets[outer_start]) + byte_end = int(inner_offsets[outer_start + 1]) + data = np.asarray( + inner_values.slice(byte_start, byte_end - byte_start) + ).tobytes() + with Image.open(BytesIO(data)) as image: + yield np.asarray(image.convert("RGB"), dtype=np.uint8) + return + + for index in range(len(values)): + if not valid[index]: + continue + outer_start = int(outer_offsets[index]) + outer_end = int(outer_offsets[index + 1]) + if outer_end <= outer_start: + continue value = values[index].as_py() if not value: - return None - return bytes(cast(bytes | bytearray | list[int], value[0])) - inner_offsets = np.asarray(inner.offsets) - byte_start = int(inner_offsets[outer_start]) - byte_end = int(inner_offsets[outer_start + 1]) - payload = inner.values.slice(byte_start, byte_end - byte_start) - return np.asarray(payload).tobytes() + continue + data = bytes(cast(bytes | bytearray | list[int], value[0])) + with Image.open(BytesIO(data)) as image: + yield np.asarray(image.convert("RGB"), dtype=np.uint8) def _is_valid(values: pa.Array | pa.ChunkedArray) -> pa.Array | pa.ChunkedArray: return pc.call_function("is_valid", [values]) -__all__ = ["RerunReader", "RerunRecording", "RerunOutputMode"] +__all__ = [ + "DEFAULT_RERUN_ACTION_PREFIX", + "DEFAULT_RERUN_CAMERA_PREFIX", + "DEFAULT_RERUN_STATE_PREFIX", + "RerunReader", + "RerunRecording", + "RerunOutputMode", +] diff --git a/src/refiner/platform/client/api.py b/src/refiner/platform/client/api.py index 6413ad59..b4f33f13 100644 --- a/src/refiner/platform/client/api.py +++ b/src/refiner/platform/client/api.py @@ -200,7 +200,7 @@ def _request_raw( query_params: dict[str, Any] | None = None, json_payload: dict[str, Any] | None = None, timeout_s: float = 10.0, - retry_attempts: int = LIFECYCLE_REQUEST_ATTEMPTS, + retry_attempts: int = 1, retry_initial_delay_s: float = LIFECYCLE_RETRY_INITIAL_DELAY_S, ) -> dict[str, Any]: resolved_path = path @@ -243,7 +243,7 @@ def _request( query_params: dict[str, Any] | None = None, json_payload: dict[str, Any] | None = None, timeout_s: float = 10.0, - retry_attempts: int = LIFECYCLE_REQUEST_ATTEMPTS, + retry_attempts: int = 1, retry_initial_delay_s: float = LIFECYCLE_RETRY_INITIAL_DELAY_S, ) -> T: response_data = self._request_raw( @@ -287,6 +287,7 @@ def verify_api_key(self, *, timeout_s: float = 10.0) -> VerifyApiKeyResponse: path="/api/me", response_type=VerifyApiKeyResponse, timeout_s=timeout_s, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def report_stage_started( @@ -299,6 +300,7 @@ def report_stage_started( method="POST", path=f"/api/jobs/{job_id}/stages/{stage_index}/start", response_type=StageLifecycleResponse, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def report_stage_finished( @@ -317,6 +319,7 @@ def report_stage_finished( path=f"/api/jobs/{job_id}/stages/{stage_index}/finish", response_type=StageLifecycleResponse, json_payload=payload, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def report_stage_heartbeat( @@ -330,6 +333,7 @@ def report_stage_heartbeat( path=f"/api/jobs/{job_id}/stages/{stage_index}/heartbeat", response_type=StageLifecycleResponse, json_payload={}, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cloud_submit_job( @@ -358,6 +362,7 @@ def cloud_create_file_upload_urls( "object_ttl_secs": object_ttl_secs, }, timeout_s=30.0, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cloud_upload_file( @@ -405,6 +410,7 @@ def cloud_complete_files( "object_ttl_secs": object_ttl_secs, }, timeout_s=30.0, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_list_jobs( @@ -426,13 +432,22 @@ def cli_list_jobs( "limit": limit, "cursor": cursor, }, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_get_job(self, *, job_id: str) -> dict[str, Any]: - return self._request_raw(method="GET", path=f"/api/cli/jobs/{job_id}") + return self._request_raw( + method="GET", + path=f"/api/cli/jobs/{job_id}", + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, + ) def cli_get_job_manifest(self, *, job_id: str) -> dict[str, Any]: - return self._request_raw(method="GET", path=f"/api/cli/jobs/{job_id}/manifest") + return self._request_raw( + method="GET", + path=f"/api/cli/jobs/{job_id}/manifest", + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, + ) def cli_get_job_workers( self, @@ -450,6 +465,7 @@ def cli_get_job_workers( "limit": limit, "cursor": cursor, }, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_get_job_logs( @@ -485,6 +501,7 @@ def cli_get_job_logs( "search": search, }, timeout_s=30.0, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_get_job_metrics( @@ -509,6 +526,7 @@ def cli_get_job_metrics( "workerIds": worker_ids, }, timeout_s=30.0, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_get_job_step_metrics( @@ -533,16 +551,22 @@ def cli_get_job_step_metrics( "sort": sort, }, timeout_s=30.0, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_cancel_job(self, *, job_id: str) -> dict[str, Any]: - return self._request_raw(method="POST", path=f"/api/cli/jobs/{job_id}/cancel") + return self._request_raw( + method="POST", + path=f"/api/cli/jobs/{job_id}/cancel", + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, + ) def cli_list_secrets(self, *, env: str | None = None) -> dict[str, Any]: return self._request_raw( method="GET", path="/api/cli/secrets", query_params={"env": env}, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_set_secret( @@ -552,6 +576,7 @@ def cli_set_secret( method="POST", path="/api/cli/secrets", json_payload={"env": env, "name": name, "value": value}, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def cli_delete_secret(self, *, name: str, env: str = "default") -> dict[str, Any]: @@ -559,6 +584,7 @@ def cli_delete_secret(self, *, name: str, env: str = "default") -> dict[str, Any method="DELETE", path=f"/api/cli/secrets/{quote(name, safe='')}", query_params={"env": env}, + retry_attempts=LIFECYCLE_REQUEST_ATTEMPTS, ) def start_worker_services( diff --git a/src/refiner/platform/manifest.py b/src/refiner/platform/manifest.py index 439f96bc..ac9635a1 100644 --- a/src/refiner/platform/manifest.py +++ b/src/refiner/platform/manifest.py @@ -14,14 +14,14 @@ from pathlib import Path from typing import TYPE_CHECKING, Any +from packaging.requirements import InvalidRequirement, Requirement +from refiner.pipeline.builtins import iter_pipeline_builtin_specs + if TYPE_CHECKING: from refiner.pipeline.planning import PlannedStage -from packaging.requirements import InvalidRequirement, Requirement - _REDACTION_PLACEHOLDER = "REDACTED_SECRET" _NORMALIZED_DEPENDENCY_SEPARATOR_PATTERN = re.compile(r"[-_.]+") -_REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__" def _redact_captured_text(text: str, *, secret_values: Sequence[str]) -> str: @@ -252,16 +252,8 @@ def build_run_manifest( for stage in pipeline_stages or (): pipeline = stage.pipeline pipeline_refiner_extras.update(pipeline.source.required_refiner_extras()) - for step in pipeline.pipeline_steps: - for candidate in getattr(step, "ops", (step,)): - for attr in ("fn", "predicate"): - spec = getattr( - getattr(candidate, attr, None), _REFINER_BUILTIN_CALL_ATTR, None - ) - if isinstance(spec, dict): - declared = spec.get("refiner_extras", ()) - if isinstance(declared, tuple): - pipeline_refiner_extras.update(declared) + for spec in iter_pipeline_builtin_specs(pipeline): + pipeline_refiner_extras.update(spec.refiner_extras) if pipeline.sink is not None: pipeline_refiner_extras.update(pipeline.sink.required_refiner_extras()) if isinstance(refiner_extras, str): diff --git a/src/refiner/services/discovery.py b/src/refiner/services/discovery.py index 966fae1f..1bb71734 100644 --- a/src/refiner/services/discovery.py +++ b/src/refiner/services/discovery.py @@ -4,69 +4,26 @@ from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Any -from refiner.pipeline.steps import ( - FnAsyncRowStep, - FnBatchStep, - FnFlatMapStep, - FnRowStep, - FnTableStep, -) +from refiner.pipeline.builtins import iter_pipeline_builtin_specs from refiner.services.base import RuntimeServiceSpec if TYPE_CHECKING: from refiner.pipeline import RefinerPipeline -_REFINER_BUILTIN_CALL_ATTR = "__refiner_builtin_call__" - - -def _builtin_description(fn: Any) -> dict[str, Any] | None: - spec = getattr(fn, _REFINER_BUILTIN_CALL_ATTR, None) - if not isinstance(spec, dict): - return None - name = spec.get("name") - if not isinstance(name, str) or not name: - return None - args = spec.get("args") - if not isinstance(args, dict): - return None - services = spec.get("services", ()) - if not isinstance(services, (list, tuple)): - return None - parsed_services: list[RuntimeServiceSpec] = [] - for service in services: - if not isinstance(service, RuntimeServiceSpec): - return None - parsed_services.append(service) - return {"name": name, "args": args, "services": tuple(parsed_services)} - - def collect_pipeline_services( pipeline: "RefinerPipeline", ) -> tuple[RuntimeServiceSpec, ...]: services_by_key: dict[tuple[str, str, str], RuntimeServiceSpec] = {} - for step in pipeline.pipeline_steps: - candidates: list[Any] = [] - if isinstance( - step, - FnRowStep | FnAsyncRowStep | FnBatchStep | FnFlatMapStep | FnTableStep, - ): - candidates.append(step.fn) - elif (fn := getattr(step, "fn", None)) is not None: - candidates.append(fn) - - for candidate in candidates: - builtin = _builtin_description(candidate) - if builtin is None: - continue - for service in builtin["services"]: - key = ( - service.name, - service.kind, - _service_config_key(service.config), - ) - services_by_key.setdefault(key, service) + for spec in iter_pipeline_builtin_specs(pipeline): + for service in spec.services: + key = ( + service.name, + service.kind, + _service_config_key(service.config), + ) + services_by_key.setdefault(key, service) return tuple(services_by_key.values()) diff --git a/tests/launchers/test_cloud_launcher.py b/tests/launchers/test_cloud_launcher.py index 6bf17009..ec7a0204 100644 --- a/tests/launchers/test_cloud_launcher.py +++ b/tests/launchers/test_cloud_launcher.py @@ -348,6 +348,48 @@ def manifest(**kwargs): assert captured_manifest_kwargs["refiner_extras"] == ["hf", "video"] +def test_pipeline_launch_cloud_accepts_extra_dependencies_alias( + monkeypatch, +) -> None: + captured_manifest_kwargs = {} + + def manifest(**kwargs): + captured_manifest_kwargs.update(kwargs) + return {"version": 1} + + _stub_cloud_submit(monkeypatch, manifest=manifest) + monkeypatch.setattr( + "refiner.launchers.cloud.refiner_ref_exists_on_remote", + lambda ref: True, + ) + + read_jsonl("input.jsonl").launch_cloud( + name="demo cloud", + extra_dependencies=["torch"], + ) + + assert captured_manifest_kwargs["dependencies"] == ["torch"] + + +def test_pipeline_launch_cloud_rejects_both_dependency_names() -> None: + with pytest.raises(ValueError, match="dependencies or extra_dependencies"): + read_jsonl("input.jsonl").launch_cloud( + name="demo cloud", + dependencies=["torch"], + extra_dependencies=["numpy"], + ) + + +def test_cloud_launcher_accepts_extra_dependencies_alias() -> None: + launcher = CloudLauncher( + pipeline=read_jsonl("input.jsonl"), + name="demo cloud", + extra_dependencies=["torch"], + ) + + assert launcher.dependencies == ["torch"] + + def test_pipeline_launch_cloud_passes_pipeline_stages_to_manifest( monkeypatch, ) -> None: diff --git a/tests/platform/test_client.py b/tests/platform/test_client.py index 78dd0ce5..4261e810 100644 --- a/tests/platform/test_client.py +++ b/tests/platform/test_client.py @@ -257,3 +257,55 @@ def fake_request_json(**kwargs: object) -> dict[str, object]: assert captured["method"] == "DELETE" assert captured["path"] == "/api/cli/secrets/HF%2FTOKEN?env=production" + + +@pytest.mark.parametrize( + ("operation", "expected_path"), + [ + ( + lambda client: client.cli_cancel_job(job_id="job-1"), + "/api/cli/jobs/job-1/cancel", + ), + ( + lambda client: client.cli_set_secret( + name="HF_TOKEN", + value="secret", + env="production", + ), + "/api/cli/secrets", + ), + ( + lambda client: client.cli_delete_secret( + name="HF_TOKEN", + env="production", + ), + "/api/cli/secrets/HF_TOKEN?env=production", + ), + ], +) +def test_cli_mutations_retry_transient_failures( + monkeypatch, + operation, + expected_path: str, +) -> None: + calls = 0 + captured: dict[str, object] = {} + sleeps: list[float] = [] + + def fake_request_json(**kwargs: object) -> dict[str, object]: + nonlocal calls + calls += 1 + captured.update(kwargs) + if calls == 1: + raise MacrodataApiError(503, "try again") + return {"success": True} + + monkeypatch.setattr("refiner.platform.client.api.request_json", fake_request_json) + monkeypatch.setattr("refiner.platform.client.api.time.sleep", sleeps.append) + + client = MacrodataClient(api_key="md_test", base_url="https://example.com") + assert operation(client) == {"success": True} + + assert calls == 2 + assert captured["path"] == expected_path + assert sleeps == [0.25] diff --git a/tests/platform/test_client_create_job.py b/tests/platform/test_client_create_job.py index c0205c9d..ee0c0bdf 100644 --- a/tests/platform/test_client_create_job.py +++ b/tests/platform/test_client_create_job.py @@ -2,7 +2,9 @@ from typing import cast -from refiner.platform.client import MacrodataClient +import pytest + +from refiner.platform.client import MacrodataApiError, MacrodataClient def _job_submit_response() -> dict[str, object]: @@ -47,3 +49,28 @@ def fake_request_json(**kwargs: object) -> dict[str, object]: "refiner_ref": "abc123def456", }, } + + +def test_create_job_does_not_retry_request_timeout(monkeypatch) -> None: + calls = 0 + sleeps: list[float] = [] + + def fake_request_json(**kwargs: object) -> dict[str, object]: + nonlocal calls + del kwargs + calls += 1 + raise MacrodataApiError(0, "read timed out") + + monkeypatch.setattr("refiner.platform.client.api.request_json", fake_request_json) + monkeypatch.setattr("refiner.platform.client.api.time.sleep", sleeps.append) + + client = MacrodataClient(api_key="ing_test", base_url="https://example.com") + with pytest.raises(MacrodataApiError, match="read timed out"): + client.create_job( + name="local job", + plan={"stages": [{"name": "stage_0", "steps": []}]}, + manifest={"version": 1}, + ) + + assert calls == 1 + assert sleeps == [] diff --git a/tests/platform/test_cloud_client.py b/tests/platform/test_cloud_client.py index 5027666b..6b7c8c7c 100644 --- a/tests/platform/test_cloud_client.py +++ b/tests/platform/test_cloud_client.py @@ -5,6 +5,7 @@ import httpx import msgspec +import pytest from refiner.pipeline.resources import GPU from refiner.platform.client import ( @@ -135,6 +136,29 @@ def test_cloud_client_cloud_submit_job_requires_job_and_stage_ids(monkeypatch) - raise AssertionError("expected MacrodataApiError") +def test_cloud_client_cloud_submit_job_does_not_retry_request_timeout( + monkeypatch, +) -> None: + calls = 0 + sleeps: list[float] = [] + + def fake_request_json(**kwargs: object) -> dict[str, object]: + nonlocal calls + del kwargs + calls += 1 + raise MacrodataApiError(0, "read timed out") + + monkeypatch.setattr("refiner.platform.client.api.request_json", fake_request_json) + monkeypatch.setattr("refiner.platform.client.api.time.sleep", sleeps.append) + + client = MacrodataClient(api_key="md_test", base_url="https://example.com") + with pytest.raises(MacrodataApiError, match="read timed out"): + client.cloud_submit_job(request=_request()) + + assert calls == 1 + assert sleeps == [] + + def test_cloud_client_cloud_submit_job_posts_continue_metadata(monkeypatch) -> None: captured: dict[str, object] = {} @@ -248,6 +272,29 @@ def fake_request_json(**kwargs: object) -> dict[str, object]: assert response.files[0].required_headers["x-amz-checksum-sha256"] == "checksum" +def test_cloud_client_retries_file_upload_url_creation(monkeypatch) -> None: + calls = 0 + sleeps: list[float] = [] + + def fake_request_json(**kwargs: object) -> dict[str, object]: + nonlocal calls + del kwargs + calls += 1 + if calls == 1: + raise MacrodataApiError(503, "try again") + return {"files": []} + + monkeypatch.setattr("refiner.platform.client.api.request_json", fake_request_json) + monkeypatch.setattr("refiner.platform.client.api.time.sleep", sleeps.append) + + client = MacrodataClient(api_key="md_test", base_url="https://example.com") + response = client.cloud_create_file_upload_urls(files=[]) + + assert response.files == [] + assert calls == 2 + assert sleeps == [0.25] + + def test_cloud_file_upload_status_serializes_as_wire_literal() -> None: assert msgspec.json.encode(CloudFileUploadStatus.NEW) == b'"new"' assert msgspec.json.encode(CloudFileUploadStatus.EXISTS) == b'"exists"' @@ -381,3 +428,26 @@ def fake_request_json(**kwargs: object) -> dict[str, object]: assert response.files[0].file_id == "00000000-0000-7000-8000-000000000123" assert response.files[0].uploaded_at == _TEST_TIMESTAMP assert response.files[0].expires_at is None + + +def test_cloud_client_retries_cloud_file_completion(monkeypatch) -> None: + calls = 0 + sleeps: list[float] = [] + + def fake_request_json(**kwargs: object) -> dict[str, object]: + nonlocal calls + del kwargs + calls += 1 + if calls == 1: + raise MacrodataApiError(429, "rate limited") + return {"files": []} + + monkeypatch.setattr("refiner.platform.client.api.request_json", fake_request_json) + monkeypatch.setattr("refiner.platform.client.api.time.sleep", sleeps.append) + + client = MacrodataClient(api_key="md_test", base_url="https://example.com") + response = client.cloud_complete_files(files=[]) + + assert response.files == [] + assert calls == 2 + assert sleeps == [0.25] diff --git a/tests/platform/test_manifest.py b/tests/platform/test_manifest.py index 30e9e1ba..ed456f1a 100644 --- a/tests/platform/test_manifest.py +++ b/tests/platform/test_manifest.py @@ -45,6 +45,18 @@ def _declared_refiner_extras(self) -> tuple[str, ...]: return ("zarr",) +class _OutputExtras: + def required_refiner_extras(self) -> tuple[str, ...]: + return ("s3",) + + +class _OutputExtrasSink(BaseSink): + output = _OutputExtras() + + def write_shard_block(self, shard_id: str, block: Block) -> None: + del shard_id, block + + def test_build_run_manifest_captures_script_from_argv( monkeypatch, tmp_path: Path ) -> None: @@ -256,6 +268,29 @@ def passthrough_table(table): ] +def test_build_run_manifest_preserves_custom_sink_output_extras( + monkeypatch, + tmp_path: Path, +) -> None: + script_path = tmp_path / "demo_job.py" + script_path.write_text("print('hello')\n", encoding="utf-8") + monkeypatch.setattr(sys, "argv", [str(script_path)]) + + pipeline = RefinerPipeline(_RefinerExtrasSource()).with_sink(_OutputExtrasSink()) + stages = [ + PlannedStage( + index=0, + name="stage_0", + pipeline=pipeline, + compute=StageComputeRequirements(num_workers=1), + ) + ] + + manifest = build_run_manifest(capture_dependencies=False, pipeline_stages=stages) + + assert manifest["environment"]["refiner_extras"] == ["hf", "s3"] + + def test_build_run_manifest_normalizes_refiner_extra_names( monkeypatch, tmp_path: Path ) -> None: diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index b9adee6f..7b976269 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -4,6 +4,7 @@ from pathlib import Path from typing import Any, cast +import fsspec import numpy as np import pytest @@ -12,7 +13,7 @@ from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table -from refiner.pipeline.sources.readers.rerun import RerunRecording +from refiner.pipeline.sources.readers.rerun import RerunReader, RerunRecording pytest.importorskip("rerun") @@ -146,6 +147,29 @@ def _reserved_video_name_rrd(path: Path) -> None: ) +def _colliding_camera_names_rrd(path: Path) -> None: + import rerun as rr + from PIL import Image + + rr.init("refiner_rerun_colliding_camera_names_test", recording_id="episode-cameras") + rr.save(path) + frames = np.arange(1) + blobs: list[bytes] = [] + for color in ((1, 2, 3), (4, 5, 6)): + out = BytesIO() + Image.new("RGB", (1, 1), color=color).save(out, format="PNG") + blobs.append(out.getvalue()) + for entity_path, blob in zip(("/cam/a.b", "/cam/a/b"), blobs, strict=True): + rr.send_columns( + entity_path, + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.EncodedImage.columns( + blob=[blob], + media_type=["image/png"], + ), + ) + + def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -158,9 +182,26 @@ def test_read_rerun_recording_preserves_timeline_table(tmp_path: Path) -> None: assert row["episode_id"] == "episode-a" assert list(recording.tables) == ["frame"] assert recording.tables["frame"].num_rows == 3 + assert recording.source_path == str(rrd) assert row["file_path"] == str(rrd) +def test_read_rerun_recording_can_project_before_arrow_conversion( + tmp_path: Path, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + + selected = ( + mdr.read_rerun(str(rrd), timelines=("frame",)).select("episode_id").take(1)[0] + ) + dropped = mdr.read_rerun(str(rrd), timelines=("frame",)).drop("rerun").take(1)[0] + + assert selected.to_dict() == {"episode_id": "episode-a"} + assert dropped["episode_id"] == "episode-a" + assert "rerun" not in dropped + + def test_read_rerun_recording_rejects_reserved_file_path_column( tmp_path: Path, ) -> None: @@ -203,6 +244,7 @@ def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> assert recording.static is None assert recording.source_file is not None assert recording.timelines == ("frame",) + assert recording.use_source_chunks is False def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: @@ -220,6 +262,21 @@ def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: match="include_recording=False is only supported for robotics output", ): mdr.read_rerun(str(rrd), output="recording", include_recording=False) + with pytest.raises( + ValueError, + match="Rerun recording output does not use robotics options: primary_timeline", + ): + mdr.read_rerun(str(rrd), output="recording", primary_timeline="frame") + with pytest.raises( + ValueError, + match="Rerun recording output does not use robotics options: actions, fps", + ): + mdr.read_rerun( + str(rrd), + output="recording", + actions=("/action/x",), + fps=30.0, + ) def test_read_rerun_recording_without_materialized_tables_skips_server( @@ -251,6 +308,36 @@ def test_read_rerun_recording_without_materialized_tables_skips_server( assert row["rerun"].tables == {} +def test_read_rerun_batches_small_files_in_one_staged_read( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + first = tmp_path / "first.rrd" + second = tmp_path / "second.rrd" + first.write_bytes(b"first") + second.write_bytes(b"second") + reader = RerunReader( + [str(first), str(second)], + target_shard_bytes=1024 * 1024, + ) + batch_sizes: list[int] = [] + + def fake_read_files(self: RerunReader, local_files: Any) -> Any: + del self + local_files = tuple(local_files) + batch_sizes.append(len(local_files)) + for source, local_path, _local_rrd in local_files: + assert local_path.exists() + yield DictRow({"source": source.abs_path()}) + + monkeypatch.setattr(RerunReader, "_read_files", fake_read_files) + + rows = cast(list[Row], list(reader.read_shard(reader.list_shards()[0]))) + + assert batch_sizes == [2] + assert [row["source"] for row in rows] == [str(first), str(second)] + + def test_read_rerun_robotics_mode_converts_to_robot_row(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) @@ -353,6 +440,24 @@ def test_read_rerun_robotics_rejects_reserved_implicit_video_name( ).take(1) +def test_read_rerun_robotics_rejects_derived_video_name_collision( + tmp_path: Path, +) -> None: + rrd = tmp_path / "colliding-cameras.rrd" + _colliding_camera_names_rrd(rrd) + + with pytest.raises( + ValueError, + match="derive the same output video name", + ): + mdr.read_rerun( + str(rrd), + output="robotics", + camera_prefix="/cam", + fps=30.0, + ).take(1) + + def test_read_rerun_robotics_mode_respects_explicit_selections( tmp_path: Path, ) -> None: @@ -497,24 +602,80 @@ def test_write_rerun_roundtrips_recording_row(tmp_path: Path) -> None: def test_write_rerun_uses_source_chunks_without_materialized_tables( tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, ) -> None: source = tmp_path / "tiny.rrd" output = tmp_path / "out-raw-copy" _tiny_rrd(source) - row = cast( - Any, - next( - mdr.read_rerun( - str(source), - timelines=("frame",), - materialize_tables=False, - ).source.read() + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + row = block[0] + recording = row["rerun"] + assert recording.use_source_chunks is True + assert recording.local_source is not None + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.LocalRrd", + lambda *args, **kwargs: pytest.fail( + "writer should reuse a live reader-staged RRD path" + ), + ) + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + copied = cast( + Any, next(mdr.read_rerun(str(written[0]), timelines=("frame",)).source.read()) + ) + + assert copied["episode_id"] == "episode-a" + assert copied["rerun"].tables["frame"].num_rows == 3 + + +def test_write_rerun_reuses_reader_staged_remote_source( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-remote-raw-copy" + _tiny_rrd(source) + + remote_fs = fsspec.filesystem("memory") + remote_path = f"/refiner-rerun-test/{tmp_path.name}/tiny.rrd" + remote_fs.pipe_file(remote_path, source.read_bytes()) + + source_iter = mdr.read_rerun( + (remote_path, remote_fs), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + row = block[0] + recording = row["rerun"] + assert not recording.source_file.is_local + assert recording.local_source is not None + assert recording.local_source.path is not None + assert recording.local_source.path.exists() + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.LocalRrd", + lambda *args, **kwargs: pytest.fail( + "writer should reuse the reader-staged remote RRD path" ), ) sink = RerunSink(str(output)) - sink.write_shard_block("shard-a", [row]) + sink.write_shard_block("shard-a", block) sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + assert recording.local_source.path is None written = sorted(output.glob("**/*.rrd")) assert len(written) == 1 @@ -526,6 +687,29 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( assert copied["rerun"].tables["frame"].num_rows == 3 +def test_write_rerun_rejects_timeline_filtered_metadata_only_recording( + tmp_path: Path, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-timeline-filtered-metadata-only" + _tiny_rrd(source) + + row = cast( + Any, + next( + mdr.read_rerun( + str(source), + timelines=("frame",), + materialize_tables=False, + ).source.read() + ), + ) + sink = RerunSink(str(output)) + + with pytest.raises(ValueError, match="without materialized Rerun table columns"): + sink.write_shard_block("shard-a", [row]) + + def test_write_rerun_rejects_segment_id_path_separator_in_filename( tmp_path: Path, ) -> None: From ad1b9e65c24140ccd7f13ba16290c3b0e68d70b0 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 10:52:56 +0200 Subject: [PATCH 31/65] Fast-path raw Rerun copies --- docs/writing-data/rerun.md | 9 ++-- src/refiner/pipeline/_rerun_io.py | 1 + src/refiner/pipeline/sinks/rerun.py | 14 ++++++ src/refiner/pipeline/sources/readers/rerun.py | 9 +++- tests/readers/test_rerun_reader.py | 48 +++++++++++++++++++ 5 files changed, 76 insertions(+), 5 deletions(-) diff --git a/docs/writing-data/rerun.md b/docs/writing-data/rerun.md index f68fe985..566b3c75 100644 --- a/docs/writing-data/rerun.md +++ b/docs/writing-data/rerun.md @@ -40,10 +40,11 @@ writing. ## Writer strategy -When the input row came from `read_rerun`, the writer uses Rerun's raw -`LazyChunkStream` path and writes the selected source chunks directly. This -preserves Rerun chunk metadata and avoids re-emitting large Arrow tables through -Python. +When the input row came from `read_rerun`, the writer uses the source RRD +instead of re-emitting large Arrow tables through Python. Unfiltered +single-recording copies are written as a byte-for-byte copy. Filtered writes and +multi-recording sources use Rerun's raw `LazyChunkStream` path to write the +selected source chunks directly. For pure copy jobs, use `read_rerun(..., materialize_tables=False)` before `write_rerun(...)` to skip timeline/static table materialization while keeping diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 79e7ea7c..5388c364 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -76,6 +76,7 @@ class RerunRecording: timelines: tuple[str, ...] | None = None include_static: bool = True use_source_chunks: bool = True + source_recording_count: int | None = None __all__ = ["LocalRrd", "RerunRecording"] diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index 036191ab..c65fff61 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -1,6 +1,7 @@ from __future__ import annotations import os +import shutil import tempfile import warnings from pathlib import Path @@ -173,6 +174,10 @@ def _write_source_chunks_from_path( local_path: Path, application_id: str, ) -> None: + if _can_copy_source_rrd(recording): + shutil.copyfile(local_path, path) + return + import rerun as rr with warnings.catch_warnings(): @@ -200,6 +205,15 @@ def _write_source_chunks_from_path( ) +def _can_copy_source_rrd(recording: RerunRecording) -> bool: + return ( + recording.source_recording_count == 1 + and recording.contents is None + and recording.timelines is None + and recording.include_static + ) + + def _filter_timelines( stream: Any, *, diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 00d8d33f..5ddd95fb 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -348,6 +348,7 @@ def _read_dataset( if self.output == "recording" or self.include_recording else [] ) + source_recording_count = len(store_entries) if store_entries else None entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} timelines = self.timelines if timelines is None: @@ -397,6 +398,7 @@ def _read_dataset( static=Tabular(static) if static is not None else None, application_id=application_id, recording_id=recording_id, + source_recording_count=source_recording_count, ) def _read_metadata_only_recording_rows( @@ -406,7 +408,9 @@ def _read_metadata_only_recording_rows( local_source: LocalRrd, ) -> list[DictRow]: rows = [] - for store in _recording_entries(local_path): + store_entries = _recording_entries(local_path) + source_recording_count = len(store_entries) + for store in store_entries: recording_id = str(store.recording_id) rows.append( self._recording_row( @@ -417,6 +421,7 @@ def _read_metadata_only_recording_rows( static=None, application_id=store.application_id, recording_id=recording_id, + source_recording_count=source_recording_count, ) ) return rows @@ -431,6 +436,7 @@ def _recording_row( static: Tabular | None, application_id: str | None, recording_id: str | None, + source_recording_count: int | None, ) -> DictRow: data: dict[str, Any] = { "episode_id": segment_id, @@ -445,6 +451,7 @@ def _recording_row( ), application_id=application_id, recording_id=recording_id, + source_recording_count=source_recording_count, contents=self.contents, timelines=self.timelines, include_static=self.include_static, diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 7b976269..c816df9d 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -1,5 +1,6 @@ from __future__ import annotations +from dataclasses import replace from io import BytesIO from pathlib import Path from typing import Any, cast @@ -617,6 +618,7 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( recording = row["rerun"] assert recording.use_source_chunks is True assert recording.local_source is not None + assert recording.source_recording_count == 1 monkeypatch.setattr( "refiner.pipeline.sinks.rerun.LocalRrd", @@ -624,6 +626,12 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( "writer should reuse a live reader-staged RRD path" ), ) + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun._matching_store", + lambda *args, **kwargs: pytest.fail( + "unfiltered single-recording copies should not rewrite RRD chunks" + ), + ) sink = RerunSink(str(output)) sink.write_shard_block("shard-a", block) sink.on_shard_complete("shard-a") @@ -663,6 +671,7 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert recording.local_source is not None assert recording.local_source.path is not None assert recording.local_source.path.exists() + assert recording.source_recording_count == 1 monkeypatch.setattr( "refiner.pipeline.sinks.rerun.LocalRrd", @@ -677,6 +686,45 @@ def test_write_rerun_reuses_reader_staged_remote_source( next(source_iter) assert recording.local_source.path is None + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + assert written[0].read_bytes() == remote_fs.cat(remote_path) + copied = cast( + Any, next(mdr.read_rerun(str(written[0]), timelines=("frame",)).source.read()) + ) + + assert copied["episode_id"] == "episode-a" + assert copied["rerun"].tables["frame"].num_rows == 3 + + +def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-not-direct-copy" + _tiny_rrd(source) + + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + row = block[0] + recording = replace(row["rerun"], source_recording_count=2) + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.shutil.copyfile", + lambda *args, **kwargs: pytest.fail( + "multi-recording source rows must use the chunk-selection path" + ), + ) + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", [row.update({"rerun": recording})]) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + written = sorted(output.glob("**/*.rrd")) assert len(written) == 1 copied = cast( From f48d964259153334e93c63a7d97ed59a63d030d2 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 11:01:45 +0200 Subject: [PATCH 32/65] Add local Rerun copy benchmark --- benchmark/rerun/README.md | 12 ++ benchmark/rerun/run_local_benchmark.py | 195 +++++++++++++++++++++++++ 2 files changed, 207 insertions(+) create mode 100644 benchmark/rerun/run_local_benchmark.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 3fd691ab..52e36070 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -9,6 +9,8 @@ This folder contains cloud benchmark harnesses for Rerun RRD workloads. prints case-level and stage-level timing deltas. - `refresh_aws_secrets.py`: copies short-lived credentials from an AWS CLI profile into the Macrodata workspace secret environment used by cloud jobs. +- `run_local_benchmark.py`: runs a local single-recording RRD copy benchmark + that compares the direct byte-copy path with the chunk-selection fallback. The default inputs are the ten base RRD files from: @@ -68,6 +70,16 @@ Useful options: By default the harness records the failed case and stops, so bad credentials or setup failures do not create a misleading benchmark session. +For a local smoke benchmark that does not require cloud credentials: + +```bash +uv run python benchmark/rerun/run_local_benchmark.py +``` + +The local benchmark generates a synthetic single-recording RRD, then measures +the direct-copy branch against the chunk-selection fallback on the same source +file. + Artifacts are written under `benchmark/rerun/artifacts/` by default: - one per-case result JSON diff --git a/benchmark/rerun/run_local_benchmark.py b/benchmark/rerun/run_local_benchmark.py new file mode 100644 index 00000000..e038a6e8 --- /dev/null +++ b/benchmark/rerun/run_local_benchmark.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +import argparse +import json +from contextlib import contextmanager +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from time import perf_counter +from typing import Iterator, cast + +import numpy as np + +import refiner as mdr +from refiner.pipeline.data.row import Row +from refiner.pipeline.sinks import rerun as rerun_sink + +DEFAULT_ARTIFACTS_DIR = Path(__file__).resolve().parent / "artifacts" + + +@dataclass(slots=True) +class CaseResult: + mode: str + wall_time_s: float + output_size_bytes: int + output_matches_input: bool + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Run a local Rerun copy benchmark comparing direct byte copies " + "with the chunk-selection fallback." + ) + ) + parser.add_argument("--iterations", type=int, default=3) + parser.add_argument("--chunks", type=int, default=100) + parser.add_argument("--rows-per-chunk", type=int, default=1000) + parser.add_argument("--artifacts-dir", type=Path, default=DEFAULT_ARTIFACTS_DIR) + parser.add_argument("--run-token") + return parser.parse_args() + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _git_ref() -> str: + import subprocess + + return subprocess.check_output( + ["git", "rev-parse", "HEAD"], + cwd=Path(__file__).resolve().parents[2], + text=True, + ).strip() + + +def _package_version(name: str) -> str: + try: + from importlib.metadata import version + + return version(name) + except Exception: + return "unknown" + + +def _package_versions() -> dict[str, str]: + return { + "macrodata-refiner": _package_version("macrodata-refiner"), + "rerun-sdk": _package_version("rerun-sdk"), + "pyarrow": _package_version("pyarrow"), + "numpy": _package_version("numpy"), + } + + +def _generate_input(path: Path, *, chunks: int, rows_per_chunk: int) -> None: + import rerun as rr + + rec = rr.RecordingStream("refiner-rerun-local-benchmark", recording_id="episode-a") + rec.save(path) + for chunk_index in range(chunks): + start = chunk_index * rows_per_chunk + frames = np.arange(start, start + rows_per_chunk, dtype=np.int64) + values = np.asarray(frames, dtype=np.float64) + rec.send_columns( + "/action/x", + indexes=[rr.TimeColumn("frame", sequence=frames)], + columns=rr.Scalars.columns(scalars=values), + ) + rec.flush() + rec.disconnect() + + +@contextmanager +def _force_chunk_fallback() -> Iterator[None]: + original = rerun_sink._can_copy_source_rrd + # This is a deliberate benchmark switch: compare the optimized path to the + # existing chunk-selection fallback on the same source file. + rerun_sink._can_copy_source_rrd = lambda recording: False # type: ignore[assignment] + try: + yield + finally: + rerun_sink._can_copy_source_rrd = original + + +def _run_copy_case(source: Path, output: Path, *, force_fallback: bool) -> CaseResult: + source_row = next( + mdr.read_rerun(str(source), materialize_tables=False).source.read() + ) + block = cast(list[Row], source_row) + sink = rerun_sink.RerunSink(str(output)) + start = perf_counter() + if force_fallback: + with _force_chunk_fallback(): + sink.write_shard_block("shard-a", block) + else: + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + wall_time_s = perf_counter() - start + written = sorted(output.glob("**/*.rrd")) + if len(written) != 1: + raise RuntimeError(f"expected one output RRD, got {len(written)}") + output_path = written[0] + return CaseResult( + mode="chunk-fallback" if force_fallback else "direct-copy", + wall_time_s=wall_time_s, + output_size_bytes=output_path.stat().st_size, + output_matches_input=output_path.read_bytes() == source.read_bytes(), + ) + + +def main() -> int: + args = _parse_args() + run_token = ( + args.run_token + or f"local-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}" + ) + artifacts_dir = args.artifacts_dir / run_token + artifacts_dir.mkdir(parents=True, exist_ok=True) + input_path = artifacts_dir / "input.rrd" + _generate_input(input_path, chunks=args.chunks, rows_per_chunk=args.rows_per_chunk) + + results: list[CaseResult] = [] + for iteration in range(args.iterations): + for mode_name, force_fallback in ( + ("direct-copy", False), + ("chunk-fallback", True), + ): + output_dir = artifacts_dir / f"{mode_name}-{iteration:02d}" + output_dir.mkdir(parents=True, exist_ok=True) + result = _run_copy_case( + input_path, output_dir, force_fallback=force_fallback + ) + results.append(result) + print( + f"{mode_name} iteration {iteration}: " + f"{result.wall_time_s:.3f}s output={result.output_size_bytes}" + ) + + summary = { + "run_token": run_token, + "git_ref": _git_ref(), + "started_at_utc": _utc_now(), + "input": { + "path": str(input_path), + "size_bytes": input_path.stat().st_size, + "chunks": args.chunks, + "rows_per_chunk": args.rows_per_chunk, + }, + "iterations": args.iterations, + "results": [asdict(result) for result in results], + "package_versions": _package_versions(), + } + summary_path = artifacts_dir / "summary.json" + summary_path.write_text( + json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8" + ) + print(f"Summary written to {summary_path}") + + direct = [result.wall_time_s for result in results if result.mode == "direct-copy"] + fallback = [ + result.wall_time_s for result in results if result.mode == "chunk-fallback" + ] + if direct and fallback: + print( + "direct-copy avg=" + f"{sum(direct) / len(direct):.3f}s " + "chunk-fallback avg=" + f"{sum(fallback) / len(fallback):.3f}s" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) From 1ceceaa046aa836088cc230cf702250f38f684f6 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 11:07:46 +0200 Subject: [PATCH 33/65] Avoid double metadata scans for raw Rerun copies --- src/refiner/pipeline/sources/readers/rerun.py | 18 +++++---- tests/readers/test_rerun_reader.py | 40 +++++++++++++++++++ 2 files changed, 50 insertions(+), 8 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 5ddd95fb..87a8b563 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -293,10 +293,11 @@ def _read_files( ) server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] for source, local_path, local_source in local_files: + store_entries = _recording_entries(local_path) rows = self._read_metadata_only_recording_rows( source, - local_path, local_source, + store_entries, ) if rows: yield from rows @@ -342,12 +343,14 @@ def _read_dataset( local_path: Path, local_source: LocalRrd, dataset: Any, + store_entries: Sequence[Any] | None = None, ) -> Iterator[SourceUnit]: - store_entries = ( - _recording_entries(local_path) - if self.output == "recording" or self.include_recording - else [] - ) + if store_entries is None: + store_entries = ( + _recording_entries(local_path) + if self.output == "recording" or self.include_recording + else [] + ) source_recording_count = len(store_entries) if store_entries else None entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} timelines = self.timelines @@ -404,11 +407,10 @@ def _read_dataset( def _read_metadata_only_recording_rows( self, source: DataFile, - local_path: Path, local_source: LocalRrd, + store_entries: Sequence[Any], ) -> list[DictRow]: rows = [] - store_entries = _recording_entries(local_path) source_recording_count = len(store_entries) for store in store_entries: recording_id = str(store.recording_id) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index c816df9d..3fc3d221 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -248,6 +248,46 @@ def test_read_rerun_recording_can_skip_table_materialization(tmp_path: Path) -> assert recording.use_source_chunks is False +def test_read_rerun_recording_without_materialized_tables_scans_metadata_once( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + rrd = tmp_path / "tiny.rrd" + _tiny_rrd(rrd) + calls = 0 + + def fake_recording_entries(*args: Any, **kwargs: Any) -> list[Any]: + nonlocal calls + del args, kwargs + calls += 1 + return [ + type( + "Store", + (), + {"recording_id": "episode-a", "application_id": "refiner"}, + )() + ] + + monkeypatch.setattr( + "refiner.pipeline.sources.readers.rerun._recording_entries", + fake_recording_entries, + ) + + row = cast( + Any, + next( + mdr.read_rerun( + str(rrd), + materialize_tables=False, + ).source.read() + ), + ) + + assert isinstance(row, list) + assert row[0]["episode_id"] == "episode-a" + assert calls == 1 + + def test_read_rerun_rejects_ignored_mode_options(tmp_path: Path) -> None: rrd = tmp_path / "tiny.rrd" _tiny_rrd(rrd) From cd3fb0557c2c5c659376b86c7676eeb8013d831a Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 11:11:13 +0200 Subject: [PATCH 34/65] Hardlink staged Rerun copies on local filesystems --- src/refiner/pipeline/sinks/rerun.py | 5 +++- tests/readers/test_rerun_reader.py | 43 +++++++++++++++++++++++++++++ 2 files changed, 47 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index c65fff61..c7180eae 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -175,7 +175,10 @@ def _write_source_chunks_from_path( application_id: str, ) -> None: if _can_copy_source_rrd(recording): - shutil.copyfile(local_path, path) + try: + os.link(local_path, path) + except OSError: + shutil.copyfile(local_path, path) return import rerun as rr diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 3fc3d221..f6832c9d 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -775,6 +775,49 @@ def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordin assert copied["rerun"].tables["frame"].num_rows == 3 +def test_write_rerun_prefers_hardlink_for_local_single_recording_copy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-hardlink-copy" + _tiny_rrd(source) + + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + row = block[0] + recording = row["rerun"] + assert recording.local_source is not None + assert recording.local_source.path is not None + staged_path = recording.local_source.path + + link_calls: list[tuple[str, str]] = [] + + def fake_link(src: str | Path, dst: str | Path) -> None: + link_calls.append((str(src), str(dst))) + Path(dst).write_bytes(Path(src).read_bytes()) + + monkeypatch.setattr("refiner.pipeline.sinks.rerun.os.link", fake_link) + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.shutil.copyfile", + lambda *args, **kwargs: pytest.fail("hardlink path should not copy bytes"), + ) + + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + assert link_calls + assert link_calls[0][0] == str(staged_path) + + def test_write_rerun_rejects_timeline_filtered_metadata_only_recording( tmp_path: Path, ) -> None: From dc8eb99351dafc73e486d8f2a4b91a824084e38b Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 11:22:54 +0200 Subject: [PATCH 35/65] Refine raw Rerun fast path and batch benchmark --- benchmark/rerun/README.md | 3 ++- benchmark/rerun/run_local_benchmark.py | 35 +++++++++++++++++++------- src/refiner/pipeline/sinks/rerun.py | 11 ++++---- tests/readers/test_rerun_reader.py | 35 ++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 16 deletions(-) diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 52e36070..d9e09764 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -78,7 +78,8 @@ uv run python benchmark/rerun/run_local_benchmark.py The local benchmark generates a synthetic single-recording RRD, then measures the direct-copy branch against the chunk-selection fallback on the same source -file. +file. Use `--writes-per-iteration` to repeat the same shard write within one +timed run when you want to amplify per-row writer overhead. Artifacts are written under `benchmark/rerun/artifacts/` by default: diff --git a/benchmark/rerun/run_local_benchmark.py b/benchmark/rerun/run_local_benchmark.py index e038a6e8..c6998bea 100644 --- a/benchmark/rerun/run_local_benchmark.py +++ b/benchmark/rerun/run_local_benchmark.py @@ -23,6 +23,7 @@ class CaseResult: mode: str wall_time_s: float output_size_bytes: int + output_file_count: int output_matches_input: bool @@ -36,6 +37,7 @@ def _parse_args() -> argparse.Namespace: parser.add_argument("--iterations", type=int, default=3) parser.add_argument("--chunks", type=int, default=100) parser.add_argument("--rows-per-chunk", type=int, default=1000) + parser.add_argument("--writes-per-iteration", type=int, default=1) parser.add_argument("--artifacts-dir", type=Path, default=DEFAULT_ARTIFACTS_DIR) parser.add_argument("--run-token") return parser.parse_args() @@ -103,7 +105,13 @@ def _force_chunk_fallback() -> Iterator[None]: rerun_sink._can_copy_source_rrd = original -def _run_copy_case(source: Path, output: Path, *, force_fallback: bool) -> CaseResult: +def _run_copy_case( + source: Path, + output: Path, + *, + force_fallback: bool, + writes_per_iteration: int, +) -> CaseResult: source_row = next( mdr.read_rerun(str(source), materialize_tables=False).source.read() ) @@ -112,20 +120,26 @@ def _run_copy_case(source: Path, output: Path, *, force_fallback: bool) -> CaseR start = perf_counter() if force_fallback: with _force_chunk_fallback(): - sink.write_shard_block("shard-a", block) + for _ in range(writes_per_iteration): + sink.write_shard_block("shard-a", block) else: - sink.write_shard_block("shard-a", block) + for _ in range(writes_per_iteration): + sink.write_shard_block("shard-a", block) sink.on_shard_complete("shard-a") wall_time_s = perf_counter() - start written = sorted(output.glob("**/*.rrd")) - if len(written) != 1: - raise RuntimeError(f"expected one output RRD, got {len(written)}") - output_path = written[0] + if len(written) != writes_per_iteration: + raise RuntimeError( + f"expected {writes_per_iteration} output RRDs, got {len(written)}" + ) return CaseResult( mode="chunk-fallback" if force_fallback else "direct-copy", wall_time_s=wall_time_s, - output_size_bytes=output_path.stat().st_size, - output_matches_input=output_path.read_bytes() == source.read_bytes(), + output_size_bytes=sum(path.stat().st_size for path in written), + output_file_count=len(written), + output_matches_input=all( + path.read_bytes() == source.read_bytes() for path in written + ), ) @@ -149,7 +163,10 @@ def main() -> int: output_dir = artifacts_dir / f"{mode_name}-{iteration:02d}" output_dir.mkdir(parents=True, exist_ok=True) result = _run_copy_case( - input_path, output_dir, force_fallback=force_fallback + input_path, + output_dir, + force_fallback=force_fallback, + writes_per_iteration=args.writes_per_iteration, ) results.append(result) print( diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index c7180eae..df9a596b 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -75,12 +75,6 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: return count def _write_recording(self, recording: RerunRecording, relpath: str) -> None: - check_required_dependencies( - "write_rerun", - [("rerun", "rerun-sdk")], - dist="rerun", - ) - def write_local(path: Path) -> None: if ( self.write_footer @@ -181,6 +175,11 @@ def _write_source_chunks_from_path( shutil.copyfile(local_path, path) return + check_required_dependencies( + "write_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) import rerun as rr with warnings.catch_warnings(): diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index f6832c9d..aec51515 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -818,6 +818,41 @@ def fake_link(src: str | Path, dst: str | Path) -> None: assert link_calls[0][0] == str(staged_path) +def test_write_rerun_direct_copy_does_not_require_rerun_sdk( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-no-rerun-sdk" + _tiny_rrd(source) + + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + row = block[0] + recording = row["rerun"] + assert recording.use_source_chunks is True + assert recording.source_recording_count == 1 + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.check_required_dependencies", + lambda *args, **kwargs: pytest.fail( + "direct-copy raw writes should not require rerun-sdk" + ), + ) + + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 1 + + def test_write_rerun_rejects_timeline_filtered_metadata_only_recording( tmp_path: Path, ) -> None: From 06c04f36413856b8a08b9619ef1adbb2bc10da89 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 11:25:37 +0200 Subject: [PATCH 36/65] Trim raw Rerun source chunk overhead --- src/refiner/pipeline/sinks/rerun.py | 2 +- tests/readers/test_rerun_reader.py | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index df9a596b..cbf3bf5c 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -141,7 +141,7 @@ def _write_source_chunks( ) -> None: local_source = recording.local_source local_source_path = local_source.path if local_source is not None else None - if local_source_path is not None and local_source_path.exists(): + if local_source_path is not None: _write_source_chunks_from_path( recording, path, diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index aec51515..a243806b 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -660,6 +660,10 @@ def test_write_rerun_uses_source_chunks_without_materialized_tables( assert recording.local_source is not None assert recording.source_recording_count == 1 + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.Path.exists", + lambda *args, **kwargs: pytest.fail("raw source chunks should not stat path"), + ) monkeypatch.setattr( "refiner.pipeline.sinks.rerun.LocalRrd", lambda *args, **kwargs: pytest.fail( From af65e540fe8e8de42fc7fb4d80e3b75169632497 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 12:28:09 +0200 Subject: [PATCH 37/65] Speed up batched Rerun writer loop --- src/refiner/pipeline/sinks/rerun.py | 92 ++++++++++++++++++++++------- tests/readers/test_rerun_reader.py | 35 +++++++++++ 2 files changed, 107 insertions(+), 20 deletions(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index cbf3bf5c..ad2ffde3 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -6,7 +6,7 @@ import warnings from pathlib import Path from string import Formatter -from typing import Any +from typing import Any, Callable import pyarrow as pa @@ -39,29 +39,32 @@ def __init__( self.output = DataFolder.resolve(output) self.filename_template = filename_template self._uses_segment_id = "segment_id" in template_fields + self._render_relpath = _compile_relpath_renderer( + filename_template, + uses_segment_id=self._uses_segment_id, + ) self.app_id = app_id self.write_footer = write_footer self._row_indices: dict[str, int] = {} self._written_relpaths: dict[str, set[str]] = {} + self._created_local_parents: set[Path] = set() def _declared_refiner_extras(self) -> tuple[str, ...]: return ("rerun",) def write_shard_block(self, shard_id: str, block: Block) -> int: count = 0 + worker_id = get_active_worker_token() + row_index = self._row_indices.get(shard_id, 0) + written_relpaths = self._written_relpaths.setdefault(shard_id, set()) for row in block: recording = _recording_from_row(row) - row_index = self._row_indices.get(shard_id, 0) - self._row_indices[shard_id] = row_index + 1 - relpath = _render_relpath( - self.filename_template, + relpath = self._render_relpath( shard_id=shard_id, - worker_id=get_active_worker_token(), + worker_id=worker_id, row_index=row_index, segment_id=recording.segment_id, - uses_segment_id=self._uses_segment_id, ) - written_relpaths = self._written_relpaths.setdefault(shard_id, set()) if relpath in written_relpaths: raise ValueError( "write_rerun filename_template rendered duplicate output path " @@ -69,6 +72,8 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: ) self._write_recording(recording, relpath) written_relpaths.add(relpath) + row_index += 1 + self._row_indices[shard_id] = row_index count += 1 if count: log_throughput("files_written", count, shard_id=shard_id, unit="files") @@ -93,7 +98,7 @@ def write_local(path: Path) -> None: target = self.output.file(relpath) if target.is_local: local_path = Path(target.abs_path()) - local_path.parent.mkdir(parents=True, exist_ok=True) + self._ensure_local_parent(local_path.parent) write_local(local_path) return @@ -125,6 +130,11 @@ def build_reducer(self) -> BaseSink | None: reducer_name="write_rerun_reduce", ) + def _ensure_local_parent(self, parent: Path) -> None: + if parent not in self._created_local_parents: + parent.mkdir(parents=True, exist_ok=True) + self._created_local_parents.add(parent) + def _recording_from_row(row: Row) -> RerunRecording: value = row.get("rerun") @@ -369,6 +379,51 @@ def _validate_filename_template(filename_template: str) -> set[str]: return fields +def _compile_relpath_renderer( + filename_template: str, + *, + uses_segment_id: bool, +) -> Callable[..., str]: + parts: list[tuple[str, str | None]] = [] + for literal_text, field_name, format_spec, conversion in Formatter().parse( + filename_template + ): + if conversion is not None or format_spec: + raise ValueError("filename_template only supports plain named fields") + parts.append((literal_text, field_name)) + + def render( + *, + shard_id: str, + worker_id: str, + row_index: int, + segment_id: str, + ) -> str: + normalized_segment_id = ( + _normalize_path_segment(segment_id, "segment_id") + if uses_segment_id + else segment_id + ) + pieces: list[str] = [] + for literal_text, field_name in parts: + pieces.append(literal_text) + if field_name is None: + continue + if field_name == "shard_id": + pieces.append(shard_id) + elif field_name == "worker_id": + pieces.append(worker_id) + elif field_name == "row_index": + pieces.append(str(row_index)) + elif field_name == "segment_id": + pieces.append(normalized_segment_id) + else: + raise AssertionError(f"unexpected filename field {field_name!r}") + return _normalize_relpath("".join(pieces), "rendered filename") + + return render + + def _render_relpath( filename_template: str, *, @@ -378,17 +433,14 @@ def _render_relpath( segment_id: str, uses_segment_id: bool, ) -> str: - field_values: dict[str, object] = { - "shard_id": shard_id, - "worker_id": worker_id, - "row_index": row_index, - "segment_id": segment_id, - } - if uses_segment_id: - field_values["segment_id"] = _normalize_path_segment(segment_id, "segment_id") - return _normalize_relpath( - filename_template.format(**field_values), - "rendered filename", + return _compile_relpath_renderer( + filename_template, + uses_segment_id=uses_segment_id, + )( + shard_id=shard_id, + worker_id=worker_id, + row_index=row_index, + segment_id=segment_id, ) diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index a243806b..fa16503a 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -857,6 +857,41 @@ def test_write_rerun_direct_copy_does_not_require_rerun_sdk( assert len(written) == 1 +def test_write_rerun_caches_local_parent_directory_creation( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + output = tmp_path / "out-mkdir-cache" + _tiny_rrd(source) + + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + + mkdir_calls: list[Path] = [] + original_mkdir = Path.mkdir + + def fake_mkdir(self: Path, *args: Any, **kwargs: Any) -> Any: + mkdir_calls.append(self) + return original_mkdir(self, *args, **kwargs) + + monkeypatch.setattr("refiner.pipeline.sinks.rerun.Path.mkdir", fake_mkdir) + + sink = RerunSink(str(output)) + sink.write_shard_block("shard-a", block) + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + + written = sorted(output.glob("**/*.rrd")) + assert len(written) == 2 + assert len(mkdir_calls) == 1 + + def test_write_rerun_rejects_timeline_filtered_metadata_only_recording( tmp_path: Path, ) -> None: From 7323294c2a26e06f220a727612a05e59d45833f6 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 12:35:08 +0200 Subject: [PATCH 38/65] Refine Rerun benchmark timing and writer loop --- benchmark/rerun/run_local_benchmark.py | 20 +++++++++------ src/refiner/pipeline/sinks/rerun.py | 35 +++++++++++++++++--------- 2 files changed, 35 insertions(+), 20 deletions(-) diff --git a/benchmark/rerun/run_local_benchmark.py b/benchmark/rerun/run_local_benchmark.py index c6998bea..ed4062fd 100644 --- a/benchmark/rerun/run_local_benchmark.py +++ b/benchmark/rerun/run_local_benchmark.py @@ -6,7 +6,7 @@ from dataclasses import asdict, dataclass from datetime import datetime, timezone from pathlib import Path -from time import perf_counter +from time import perf_counter_ns from typing import Iterator, cast import numpy as np @@ -21,11 +21,15 @@ @dataclass(slots=True) class CaseResult: mode: str - wall_time_s: float + wall_time_ns: int output_size_bytes: int output_file_count: int output_matches_input: bool + @property + def wall_time_s(self) -> float: + return self.wall_time_ns / 1_000_000_000 + def _parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( @@ -117,7 +121,7 @@ def _run_copy_case( ) block = cast(list[Row], source_row) sink = rerun_sink.RerunSink(str(output)) - start = perf_counter() + start = perf_counter_ns() if force_fallback: with _force_chunk_fallback(): for _ in range(writes_per_iteration): @@ -126,7 +130,7 @@ def _run_copy_case( for _ in range(writes_per_iteration): sink.write_shard_block("shard-a", block) sink.on_shard_complete("shard-a") - wall_time_s = perf_counter() - start + wall_time_ns = perf_counter_ns() - start written = sorted(output.glob("**/*.rrd")) if len(written) != writes_per_iteration: raise RuntimeError( @@ -134,7 +138,7 @@ def _run_copy_case( ) return CaseResult( mode="chunk-fallback" if force_fallback else "direct-copy", - wall_time_s=wall_time_s, + wall_time_ns=wall_time_ns, output_size_bytes=sum(path.stat().st_size for path in written), output_file_count=len(written), output_matches_input=all( @@ -171,7 +175,7 @@ def main() -> int: results.append(result) print( f"{mode_name} iteration {iteration}: " - f"{result.wall_time_s:.3f}s output={result.output_size_bytes}" + f"{result.wall_time_s:.6f}s output={result.output_size_bytes}" ) summary = { @@ -201,9 +205,9 @@ def main() -> int: if direct and fallback: print( "direct-copy avg=" - f"{sum(direct) / len(direct):.3f}s " + f"{sum(direct) / len(direct):.6f}s " "chunk-fallback avg=" - f"{sum(fallback) / len(fallback):.3f}s" + f"{sum(fallback) / len(fallback):.6f}s" ) return 0 diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index ad2ffde3..a836ecb0 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -37,7 +37,11 @@ def __init__( ) -> None: template_fields = _validate_filename_template(filename_template) self.output = DataFolder.resolve(output) + self._local_output_root = ( + Path(self.output.abs_path()) if self.output.is_local else None + ) self.filename_template = filename_template + self._uses_row_index = "row_index" in template_fields self._uses_segment_id = "segment_id" in template_fields self._render_relpath = _compile_relpath_renderer( filename_template, @@ -56,7 +60,11 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: count = 0 worker_id = get_active_worker_token() row_index = self._row_indices.get(shard_id, 0) - written_relpaths = self._written_relpaths.setdefault(shard_id, set()) + written_relpaths = ( + None + if self._uses_row_index + else self._written_relpaths.setdefault(shard_id, set()) + ) for row in block: recording = _recording_from_row(row) relpath = self._render_relpath( @@ -65,16 +73,17 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: row_index=row_index, segment_id=recording.segment_id, ) - if relpath in written_relpaths: - raise ValueError( - "write_rerun filename_template rendered duplicate output path " - f"{relpath!r}; include {{row_index}} or another unique row field" - ) + if written_relpaths is not None: + if relpath in written_relpaths: + raise ValueError( + "write_rerun filename_template rendered duplicate output path " + f"{relpath!r}; include {{row_index}} or another unique row field" + ) + written_relpaths.add(relpath) self._write_recording(recording, relpath) - written_relpaths.add(relpath) row_index += 1 - self._row_indices[shard_id] = row_index count += 1 + self._row_indices[shard_id] = row_index if count: log_throughput("files_written", count, shard_id=shard_id, unit="files") return count @@ -95,13 +104,14 @@ def write_local(path: Path) -> None: write_footer=self.write_footer, ) - target = self.output.file(relpath) - if target.is_local: - local_path = Path(target.abs_path()) + local_output_root = self._local_output_root + if local_output_root is not None: + local_path = local_output_root / relpath self._ensure_local_parent(local_path.parent) write_local(local_path) return + target = self.output.file(relpath) with tempfile.TemporaryDirectory(prefix="refiner-rerun-write-") as tmpdir: local_path = Path(tmpdir) / os.path.basename(relpath) write_local(local_path) @@ -109,7 +119,8 @@ def write_local(path: Path) -> None: def on_shard_complete(self, shard_id: str) -> None: self._row_indices.pop(shard_id, None) - self._written_relpaths.pop(shard_id, None) + if not self._uses_row_index: + self._written_relpaths.pop(shard_id, None) def describe(self) -> tuple[str, str, dict[str, object]]: return ( From 8ec50e2ef9669123ff118d939a4e66839b50aa0c Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 12:38:08 +0200 Subject: [PATCH 39/65] Reduce local Rerun writer path overhead --- src/refiner/pipeline/sinks/rerun.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index a836ecb0..eb211b51 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -38,7 +38,7 @@ def __init__( template_fields = _validate_filename_template(filename_template) self.output = DataFolder.resolve(output) self._local_output_root = ( - Path(self.output.abs_path()) if self.output.is_local else None + self.output.abs_path() if self.output.is_local else None ) self.filename_template = filename_template self._uses_row_index = "row_index" in template_fields @@ -51,7 +51,7 @@ def __init__( self.write_footer = write_footer self._row_indices: dict[str, int] = {} self._written_relpaths: dict[str, set[str]] = {} - self._created_local_parents: set[Path] = set() + self._created_local_parents: set[str] = set() def _declared_refiner_extras(self) -> tuple[str, ...]: return ("rerun",) @@ -89,7 +89,7 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: return count def _write_recording(self, recording: RerunRecording, relpath: str) -> None: - def write_local(path: Path) -> None: + def write_local(path: Path | str) -> None: if ( self.write_footer and recording.use_source_chunks @@ -106,8 +106,8 @@ def write_local(path: Path) -> None: local_output_root = self._local_output_root if local_output_root is not None: - local_path = local_output_root / relpath - self._ensure_local_parent(local_path.parent) + local_path = f"{local_output_root}/{relpath}" + self._ensure_local_parent(os.path.dirname(local_path)) write_local(local_path) return @@ -141,9 +141,9 @@ def build_reducer(self) -> BaseSink | None: reducer_name="write_rerun_reduce", ) - def _ensure_local_parent(self, parent: Path) -> None: + def _ensure_local_parent(self, parent: str) -> None: if parent not in self._created_local_parents: - parent.mkdir(parents=True, exist_ok=True) + Path(parent).mkdir(parents=True, exist_ok=True) self._created_local_parents.add(parent) @@ -156,7 +156,7 @@ def _recording_from_row(row: Row) -> RerunRecording: def _write_source_chunks( recording: RerunRecording, - path: Path, + path: Path | str, *, application_id: str, ) -> None: @@ -184,7 +184,7 @@ def _write_source_chunks( def _write_source_chunks_from_path( recording: RerunRecording, - path: Path, + path: Path | str, *, local_path: Path, application_id: str, @@ -290,7 +290,7 @@ def _matching_store(reader: Any, recording: RerunRecording) -> Any: def _write_recording_tables( recording: RerunRecording, - path: Path, + path: Path | str, *, application_id: str, write_footer: bool, From 9995f3b849af15bfad782536bc9627fc946c2533 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 12:43:02 +0200 Subject: [PATCH 40/65] Speed up default local Rerun writes --- src/refiner/pipeline/sinks/rerun.py | 60 ++++++++++++++++++----------- 1 file changed, 38 insertions(+), 22 deletions(-) diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index eb211b51..f9217593 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -60,29 +60,45 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: count = 0 worker_id = get_active_worker_token() row_index = self._row_indices.get(shard_id, 0) - written_relpaths = ( - None - if self._uses_row_index - else self._written_relpaths.setdefault(shard_id, set()) - ) - for row in block: - recording = _recording_from_row(row) - relpath = self._render_relpath( - shard_id=shard_id, - worker_id=worker_id, - row_index=row_index, - segment_id=recording.segment_id, + local_output_root = self._local_output_root + if ( + local_output_root is not None + and self.filename_template == _DEFAULT_FILENAME_TEMPLATE + ): + parent = f"{local_output_root}/{shard_id}__w{worker_id}" + self._ensure_local_parent(parent) + for row in block: + recording = _recording_from_row(row) + self._write_recording( + recording, + f"{parent}/{row_index}.rrd", + ) + row_index += 1 + count += 1 + else: + written_relpaths = ( + None + if self._uses_row_index + else self._written_relpaths.setdefault(shard_id, set()) ) - if written_relpaths is not None: - if relpath in written_relpaths: - raise ValueError( - "write_rerun filename_template rendered duplicate output path " - f"{relpath!r}; include {{row_index}} or another unique row field" - ) - written_relpaths.add(relpath) - self._write_recording(recording, relpath) - row_index += 1 - count += 1 + for row in block: + recording = _recording_from_row(row) + relpath = self._render_relpath( + shard_id=shard_id, + worker_id=worker_id, + row_index=row_index, + segment_id=recording.segment_id, + ) + if written_relpaths is not None: + if relpath in written_relpaths: + raise ValueError( + "write_rerun filename_template rendered duplicate output path " + f"{relpath!r}; include {{row_index}} or another unique row field" + ) + written_relpaths.add(relpath) + self._write_recording(recording, relpath) + row_index += 1 + count += 1 self._row_indices[shard_id] = row_index if count: log_throughput("files_written", count, shard_id=shard_id, unit="files") From 0d6d90ee3666e5bb73c0932937144cba29494e72 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 13:22:30 +0200 Subject: [PATCH 41/65] Speed up local-to-remote DataFile copies --- src/refiner/io/datafile.py | 11 +++++++++++ tests/io/test_datafile_datafolder.py | 21 +++++++++++++++++++++ 2 files changed, 32 insertions(+) diff --git a/src/refiner/io/datafile.py b/src/refiner/io/datafile.py index 827d2ca1..b6a7ca4a 100644 --- a/src/refiner/io/datafile.py +++ b/src/refiner/io/datafile.py @@ -146,6 +146,17 @@ def copy(self, dest: DataFileLike, *, buffer_size: int = 2 * 1024 * 1024) -> Non ): return + if self.is_local and callable(getattr(target.fs, "put_file", None)): + target.fs.makedirs(target.fs._parent(target.path), exist_ok=True) + try: + target.fs.put_file(self.abs_path(), target.path) + return + except Exception: + try: + target.fs.rm(target.path) + except FileNotFoundError: + pass + # Same-filesystem copies are usually server-side for object stores; fall back to # streaming only when the backend cannot copy directly. if self.fs is target.fs and callable(getattr(target.fs, "copy", None)): diff --git a/tests/io/test_datafile_datafolder.py b/tests/io/test_datafile_datafolder.py index e16849ad..62ad1a0a 100644 --- a/tests/io/test_datafile_datafolder.py +++ b/tests/io/test_datafile_datafolder.py @@ -85,6 +85,27 @@ def test_datafile_copy_writes_destination(tmp_path): assert dest_path.read_bytes() == b"payload" +def test_datafile_copy_uses_remote_put_file_for_local_sources(tmp_path, monkeypatch): + source_path = tmp_path / "source.txt" + source_path.write_bytes(b"payload") + fs = MemoryFileSystem() + dest = DataFile.resolve(("bucket/dest.txt", fs)) + + put_calls: list[tuple[str, str]] = [] + original_put_file = fs.put_file + + def fake_put_file(lpath, rpath, **kwargs): + put_calls.append((str(lpath), str(rpath))) + return original_put_file(lpath, rpath, **kwargs) + + monkeypatch.setattr(fs, "put_file", fake_put_file) + + DataFile.resolve(str(source_path)).copy(dest) + + assert put_calls == [(str(source_path), "bucket/dest.txt")] + assert fs.cat("bucket/dest.txt") == b"payload" + + def test_datafile_resolve_adds_hf_token_for_huggingface_http_urls(monkeypatch): captured = {} From 905221d5eadcfd417d902167d260782b1b3c8254 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 13:49:37 +0200 Subject: [PATCH 42/65] Tune local-to-remote upload buffering --- src/refiner/io/datafile.py | 6 +++++- tests/io/test_datafile_datafolder.py | 12 +++++++++--- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/src/refiner/io/datafile.py b/src/refiner/io/datafile.py index b6a7ca4a..ae1b4208 100644 --- a/src/refiner/io/datafile.py +++ b/src/refiner/io/datafile.py @@ -149,7 +149,11 @@ def copy(self, dest: DataFileLike, *, buffer_size: int = 2 * 1024 * 1024) -> Non if self.is_local and callable(getattr(target.fs, "put_file", None)): target.fs.makedirs(target.fs._parent(target.path), exist_ok=True) try: - target.fs.put_file(self.abs_path(), target.path) + target.fs.put_file( + self.abs_path(), + target.path, + block_size=8 * 1024 * 1024, + ) return except Exception: try: diff --git a/tests/io/test_datafile_datafolder.py b/tests/io/test_datafile_datafolder.py index 62ad1a0a..73a8fdeb 100644 --- a/tests/io/test_datafile_datafolder.py +++ b/tests/io/test_datafile_datafolder.py @@ -91,18 +91,24 @@ def test_datafile_copy_uses_remote_put_file_for_local_sources(tmp_path, monkeypa fs = MemoryFileSystem() dest = DataFile.resolve(("bucket/dest.txt", fs)) - put_calls: list[tuple[str, str]] = [] + put_calls: list[tuple[str, str, dict[str, Any]]] = [] original_put_file = fs.put_file def fake_put_file(lpath, rpath, **kwargs): - put_calls.append((str(lpath), str(rpath))) + put_calls.append((str(lpath), str(rpath), dict(kwargs))) return original_put_file(lpath, rpath, **kwargs) monkeypatch.setattr(fs, "put_file", fake_put_file) DataFile.resolve(str(source_path)).copy(dest) - assert put_calls == [(str(source_path), "bucket/dest.txt")] + assert put_calls == [ + ( + str(source_path), + "bucket/dest.txt", + {"block_size": 8 * 1024 * 1024}, + ) + ] assert fs.cat("bucket/dest.txt") == b"payload" From 24406de25ef801136d5d73bb0c8efa8188caae2f Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 15:49:35 +0200 Subject: [PATCH 43/65] Parallelize staged Rerun source opens --- src/refiner/pipeline/sources/readers/rerun.py | 54 ++++++++++++++----- 1 file changed, 41 insertions(+), 13 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 87a8b563..c95ea602 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -1,5 +1,6 @@ from __future__ import annotations +import concurrent.futures from collections.abc import Iterable, Iterator, Mapping, Sequence from pathlib import Path from typing import Any, Literal, cast @@ -235,7 +236,7 @@ def describe(self) -> dict[str, Any]: def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: descriptor = shard.descriptor assert isinstance(descriptor, FilePartsDescriptor) - batch: list[tuple[DataFile, Path, LocalRrd]] = [] + batch: list[tuple[DataFile, LocalRrd]] = [] batch_bytes = 0 for part in descriptor.parts: source = self.fileset.resolve_file(part.source_index, part.path) @@ -248,38 +249,35 @@ def read_shard(self, shard: Shard) -> Iterator[SourceUnit]: batch = [] batch_bytes = 0 local_source = LocalRrd(source) - try: - batch.append((source, local_source.open(), local_source)) - except BaseException: - local_source.close() - raise + batch.append((source, local_source)) batch_bytes += part_size if batch: yield from self._read_staged_batch(batch) def _read_staged_batch( self, - local_files: Sequence[tuple[DataFile, Path, LocalRrd]], + local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> Iterator[SourceUnit]: + opened_files = _open_local_sources(local_files) if self._retain_batch_local_sources(): try: - units = list(self._read_files(local_files)) + units = list(self._read_files(opened_files)) except BaseException: - _close_local_sources(local_files) + _close_local_sources(opened_files) raise if not units: - _close_local_sources(local_files) + _close_local_sources(opened_files) return try: yield cast(list[Row], units) finally: - _close_local_sources(local_files) + _close_local_sources(opened_files) return try: - yield from self._read_files(local_files) + yield from self._read_files(opened_files) finally: - _close_local_sources(local_files) + _close_local_sources(opened_files) def _read_files( self, @@ -666,6 +664,36 @@ def _close_local_sources( local_source.close() +def _open_local_sources( + local_files: Sequence[tuple[DataFile, LocalRrd]], +) -> list[tuple[DataFile, Path, LocalRrd]]: + if len(local_files) <= 1: + return [ + (source, local_source.open(), local_source) + for source, local_source in local_files + ] + + local_sources = [local_source for _source, local_source in local_files] + max_workers = min(8, len(local_files)) + try: + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: + local_paths = list(pool.map(_open_local_source, local_sources)) + except BaseException: + for local_source in local_sources: + local_source.close() + raise + return [ + (source, local_path, local_source) + for (source, local_source), local_path in zip( + local_files, local_paths, strict=True + ) + ] + + +def _open_local_source(local_source: LocalRrd) -> Path: + return local_source.open() + + def _recording_entries(local_path: Path) -> list[Any]: import rerun as rr From 0917053a33da87e9ef294ea05c197595ebd815d7 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:05:47 +0200 Subject: [PATCH 44/65] Parallelize Rerun metadata scans --- src/refiner/pipeline/sources/readers/rerun.py | 29 +++++++++++++++++-- 1 file changed, 27 insertions(+), 2 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index c95ea602..a5c6ffc9 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -290,8 +290,12 @@ def _read_files( dist="rerun", ) server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] - for source, local_path, local_source in local_files: - store_entries = _recording_entries(local_path) + for ( + source, + local_path, + local_source, + store_entries, + ) in _scan_recording_entries(local_files): rows = self._read_metadata_only_recording_rows( source, local_source, @@ -694,6 +698,27 @@ def _open_local_source(local_source: LocalRrd) -> Path: return local_source.open() +def _scan_recording_entries( + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], +) -> list[tuple[DataFile, Path, LocalRrd, list[Any]]]: + if len(local_files) <= 1: + return [ + (source, local_path, local_source, _recording_entries(local_path)) + for source, local_path, local_source in local_files + ] + + local_paths = [local_path for _source, local_path, _local_source in local_files] + max_workers = min(8, len(local_files)) + with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: + store_entries = list(pool.map(_recording_entries, local_paths)) + return [ + (source, local_path, local_source, store_entry) + for (source, local_path, local_source), store_entry in zip( + local_files, store_entries, strict=True + ) + ] + + def _recording_entries(local_path: Path) -> list[Any]: import rerun as rr From d322b0c4c3494a490d110100a879531f910d2065 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:13:44 +0200 Subject: [PATCH 45/65] Tweak Rerun batch staging concurrency --- src/refiner/pipeline/sources/readers/rerun.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index a5c6ffc9..4a55516a 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -678,7 +678,7 @@ def _open_local_sources( ] local_sources = [local_source for _source, local_source in local_files] - max_workers = min(8, len(local_files)) + max_workers = min(16, len(local_files)) try: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: local_paths = list(pool.map(_open_local_source, local_sources)) @@ -708,7 +708,7 @@ def _scan_recording_entries( ] local_paths = [local_path for _source, local_path, _local_source in local_files] - max_workers = min(8, len(local_files)) + max_workers = min(16, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: store_entries = list(pool.map(_recording_entries, local_paths)) return [ From 63305f7c90be1b22197e460e92356255bd1893d2 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:22:51 +0200 Subject: [PATCH 46/65] Revert Rerun batch staging oversubscription --- src/refiner/pipeline/sources/readers/rerun.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 4a55516a..a5c6ffc9 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -678,7 +678,7 @@ def _open_local_sources( ] local_sources = [local_source for _source, local_source in local_files] - max_workers = min(16, len(local_files)) + max_workers = min(8, len(local_files)) try: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: local_paths = list(pool.map(_open_local_source, local_sources)) @@ -708,7 +708,7 @@ def _scan_recording_entries( ] local_paths = [local_path for _source, local_path, _local_source in local_files] - max_workers = min(16, len(local_files)) + max_workers = min(8, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: store_entries = list(pool.map(_recording_entries, local_paths)) return [ From 7df6d30a6f3d925d122d14428567d0d21dd4662c Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:26:41 +0200 Subject: [PATCH 47/65] Fuse Rerun metadata staging passes --- src/refiner/pipeline/sources/readers/rerun.py | 108 ++++++++++++------ 1 file changed, 70 insertions(+), 38 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index a5c6ffc9..aaa118c7 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -258,6 +258,29 @@ def _read_staged_batch( self, local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> Iterator[SourceUnit]: + if self.output == "recording" and not self.materialize_tables: + prepared_files = _prepare_local_metadata_only_sources(local_files) + if self._retain_batch_local_sources(): + try: + units = list(self._read_metadata_only_files(prepared_files)) + except BaseException: + _close_prepared_local_sources(prepared_files) + raise + if not units: + _close_prepared_local_sources(prepared_files) + return + try: + yield cast(list[Row], units) + finally: + _close_prepared_local_sources(prepared_files) + return + + try: + yield from self._read_metadata_only_files(prepared_files) + finally: + _close_prepared_local_sources(prepared_files) + return + opened_files = _open_local_sources(local_files) if self._retain_batch_local_sources(): try: @@ -283,33 +306,30 @@ def _read_files( self, local_files: Sequence[tuple[DataFile, Path, LocalRrd]], ) -> Iterator[SourceUnit]: - if self.output == "recording" and not self.materialize_tables: - check_required_dependencies( - "read_rerun", - [("rerun", "rerun-sdk")], - dist="rerun", - ) - server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] - for ( + yield from self._read_files_with_server(local_files) + + def _read_metadata_only_files( + self, + local_files: Sequence[tuple[DataFile, Path, LocalRrd, list[Any]]], + ) -> Iterator[SourceUnit]: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] + for source, local_path, local_source, store_entries in local_files: + rows = self._read_metadata_only_recording_rows( source, - local_path, local_source, store_entries, - ) in _scan_recording_entries(local_files): - rows = self._read_metadata_only_recording_rows( - source, - local_source, - store_entries, - ) - if rows: - yield from rows - else: - server_fallback.append((source, local_path, local_source)) - if server_fallback: - yield from self._read_files_with_server(server_fallback) - return - - yield from self._read_files_with_server(local_files) + ) + if rows: + yield from rows + else: + server_fallback.append((source, local_path, local_source)) + if server_fallback: + yield from self._read_files_with_server(server_fallback) def _read_files_with_server( self, @@ -668,6 +688,13 @@ def _close_local_sources( local_source.close() +def _close_prepared_local_sources( + local_files: Iterable[tuple[DataFile, Path, LocalRrd, list[Any]]], +) -> None: + for _source, _local_path, local_source, _store_entries in local_files: + local_source.close() + + def _open_local_sources( local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> list[tuple[DataFile, Path, LocalRrd]]: @@ -698,25 +725,30 @@ def _open_local_source(local_source: LocalRrd) -> Path: return local_source.open() -def _scan_recording_entries( - local_files: Sequence[tuple[DataFile, Path, LocalRrd]], +def _prepare_local_metadata_only_sources( + local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> list[tuple[DataFile, Path, LocalRrd, list[Any]]]: if len(local_files) <= 1: - return [ - (source, local_path, local_source, _recording_entries(local_path)) - for source, local_path, local_source in local_files - ] + prepared = [] + for source, local_source in local_files: + local_path = local_source.open() + prepared.append( + (source, local_path, local_source, _recording_entries(local_path)) + ) + return prepared - local_paths = [local_path for _source, local_path, _local_source in local_files] max_workers = min(8, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: - store_entries = list(pool.map(_recording_entries, local_paths)) - return [ - (source, local_path, local_source, store_entry) - for (source, local_path, local_source), store_entry in zip( - local_files, store_entries, strict=True - ) - ] + prepared = list(pool.map(_prepare_local_metadata_only_source, local_files)) + return prepared + + +def _prepare_local_metadata_only_source( + item: tuple[DataFile, LocalRrd], +) -> tuple[DataFile, Path, LocalRrd, list[Any]]: + source, local_source = item + local_path = local_source.open() + return source, local_path, local_source, _recording_entries(local_path) def _recording_entries(local_path: Path) -> list[Any]: From 004a69c365cc0b88083978975668c43a290cb494 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:34:40 +0200 Subject: [PATCH 48/65] Revert "Fuse Rerun metadata staging passes" This reverts commit 7df6d30a6f3d925d122d14428567d0d21dd4662c. --- src/refiner/pipeline/sources/readers/rerun.py | 108 ++++++------------ 1 file changed, 38 insertions(+), 70 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index aaa118c7..a5c6ffc9 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -258,29 +258,6 @@ def _read_staged_batch( self, local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> Iterator[SourceUnit]: - if self.output == "recording" and not self.materialize_tables: - prepared_files = _prepare_local_metadata_only_sources(local_files) - if self._retain_batch_local_sources(): - try: - units = list(self._read_metadata_only_files(prepared_files)) - except BaseException: - _close_prepared_local_sources(prepared_files) - raise - if not units: - _close_prepared_local_sources(prepared_files) - return - try: - yield cast(list[Row], units) - finally: - _close_prepared_local_sources(prepared_files) - return - - try: - yield from self._read_metadata_only_files(prepared_files) - finally: - _close_prepared_local_sources(prepared_files) - return - opened_files = _open_local_sources(local_files) if self._retain_batch_local_sources(): try: @@ -306,30 +283,33 @@ def _read_files( self, local_files: Sequence[tuple[DataFile, Path, LocalRrd]], ) -> Iterator[SourceUnit]: - yield from self._read_files_with_server(local_files) - - def _read_metadata_only_files( - self, - local_files: Sequence[tuple[DataFile, Path, LocalRrd, list[Any]]], - ) -> Iterator[SourceUnit]: - check_required_dependencies( - "read_rerun", - [("rerun", "rerun-sdk")], - dist="rerun", - ) - server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] - for source, local_path, local_source, store_entries in local_files: - rows = self._read_metadata_only_recording_rows( + if self.output == "recording" and not self.materialize_tables: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + server_fallback: list[tuple[DataFile, Path, LocalRrd]] = [] + for ( source, + local_path, local_source, store_entries, - ) - if rows: - yield from rows - else: - server_fallback.append((source, local_path, local_source)) - if server_fallback: - yield from self._read_files_with_server(server_fallback) + ) in _scan_recording_entries(local_files): + rows = self._read_metadata_only_recording_rows( + source, + local_source, + store_entries, + ) + if rows: + yield from rows + else: + server_fallback.append((source, local_path, local_source)) + if server_fallback: + yield from self._read_files_with_server(server_fallback) + return + + yield from self._read_files_with_server(local_files) def _read_files_with_server( self, @@ -688,13 +668,6 @@ def _close_local_sources( local_source.close() -def _close_prepared_local_sources( - local_files: Iterable[tuple[DataFile, Path, LocalRrd, list[Any]]], -) -> None: - for _source, _local_path, local_source, _store_entries in local_files: - local_source.close() - - def _open_local_sources( local_files: Sequence[tuple[DataFile, LocalRrd]], ) -> list[tuple[DataFile, Path, LocalRrd]]: @@ -725,30 +698,25 @@ def _open_local_source(local_source: LocalRrd) -> Path: return local_source.open() -def _prepare_local_metadata_only_sources( - local_files: Sequence[tuple[DataFile, LocalRrd]], +def _scan_recording_entries( + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], ) -> list[tuple[DataFile, Path, LocalRrd, list[Any]]]: if len(local_files) <= 1: - prepared = [] - for source, local_source in local_files: - local_path = local_source.open() - prepared.append( - (source, local_path, local_source, _recording_entries(local_path)) - ) - return prepared + return [ + (source, local_path, local_source, _recording_entries(local_path)) + for source, local_path, local_source in local_files + ] + local_paths = [local_path for _source, local_path, _local_source in local_files] max_workers = min(8, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: - prepared = list(pool.map(_prepare_local_metadata_only_source, local_files)) - return prepared - - -def _prepare_local_metadata_only_source( - item: tuple[DataFile, LocalRrd], -) -> tuple[DataFile, Path, LocalRrd, list[Any]]: - source, local_source = item - local_path = local_source.open() - return source, local_path, local_source, _recording_entries(local_path) + store_entries = list(pool.map(_recording_entries, local_paths)) + return [ + (source, local_path, local_source, store_entry) + for (source, local_path, local_source), store_entry in zip( + local_files, store_entries, strict=True + ) + ] def _recording_entries(local_path: Path) -> list[Any]: From cde8671b037f4d12e6abea3b2f9c658d1142ab90 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:36:18 +0200 Subject: [PATCH 49/65] Trim Rerun metadata scan wrappers --- src/refiner/pipeline/sources/readers/rerun.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index a5c6ffc9..673de053 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -728,7 +728,15 @@ def _recording_entries(local_path: Path) -> list[Any]: "ignore", message="RRD file has no footer/manifest:.*", ) - return list(rr.experimental.RrdReader(local_path).recordings()) + reader = rr.experimental.RrdReader(local_path) + internal = getattr(reader, "_internal", None) + store_entries = ( + internal.store_entries() + if internal is not None + and callable(getattr(internal, "store_entries", None)) + else reader.recordings() + ) + return [entry for entry in store_entries if entry.kind == "recording"] except Exception as err: logger.warning( "Rerun recording metadata unavailable; falling back to server scan: {}", From 85ea27894ff6cf1eba8f21d2465497bb8327ef8a Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:45:18 +0200 Subject: [PATCH 50/65] Fast-path single-store Rerun lookup --- src/refiner/pipeline/sources/readers/rerun.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 673de053..a8d3662a 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -354,12 +354,27 @@ def _read_dataset( else [] ) source_recording_count = len(store_entries) if store_entries else None - entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} + if len(store_entries) == 1: + only_store_entry = store_entries[0] + entries_by_recording_id = None + else: + only_store_entry = None + entries_by_recording_id = { + entry.recording_id: entry for entry in store_entries + } timelines = self.timelines if timelines is None: timelines = self._timelines(dataset.schema()) for segment_id in dataset.segment_ids(): - store = entries_by_recording_id.get(segment_id) + if only_store_entry is not None: + store = ( + only_store_entry + if only_store_entry.recording_id == segment_id + else None + ) + else: + assert entries_by_recording_id is not None + store = entries_by_recording_id.get(segment_id) application_id = store.application_id if store is not None else None recording_id = store.recording_id if store is not None else segment_id view = dataset.filter_segments([segment_id]) From e2824001af81c53ea3f515b7aeacb9e07c8399c1 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 16:53:10 +0200 Subject: [PATCH 51/65] Revert "Fast-path single-store Rerun lookup" This reverts commit 85ea27894ff6cf1eba8f21d2465497bb8327ef8a. --- src/refiner/pipeline/sources/readers/rerun.py | 19 ++----------------- 1 file changed, 2 insertions(+), 17 deletions(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index a8d3662a..673de053 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -354,27 +354,12 @@ def _read_dataset( else [] ) source_recording_count = len(store_entries) if store_entries else None - if len(store_entries) == 1: - only_store_entry = store_entries[0] - entries_by_recording_id = None - else: - only_store_entry = None - entries_by_recording_id = { - entry.recording_id: entry for entry in store_entries - } + entries_by_recording_id = {entry.recording_id: entry for entry in store_entries} timelines = self.timelines if timelines is None: timelines = self._timelines(dataset.schema()) for segment_id in dataset.segment_ids(): - if only_store_entry is not None: - store = ( - only_store_entry - if only_store_entry.recording_id == segment_id - else None - ) - else: - assert entries_by_recording_id is not None - store = entries_by_recording_id.get(segment_id) + store = entries_by_recording_id.get(segment_id) application_id = store.application_id if store is not None else None recording_id = store.recording_id if store is not None else segment_id view = dataset.filter_segments([segment_id]) From 44920e21bb7dacac1092b76c88f525c6f80e77dc Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 19:57:13 +0200 Subject: [PATCH 52/65] Direct-copy raw Rerun source chunks --- benchmark/rerun/README.md | 4 ++ benchmark/rerun/compare_results.py | 35 +++++++++++ benchmark/rerun/run_cloud_benchmark.py | 16 ++++- src/refiner/pipeline/sinks/rerun.py | 11 +++- src/refiner/platform/manifest.py | 16 +++++ tests/benchmark/test_compare_results.py | 80 +++++++++++++++++++++++++ tests/platform/test_manifest.py | 33 ++++++++++ tests/readers/test_rerun_reader.py | 44 ++++++++++++++ 8 files changed, 237 insertions(+), 2 deletions(-) create mode 100644 tests/benchmark/test_compare_results.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index d9e09764..595e86c1 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -81,6 +81,10 @@ the direct-copy branch against the chunk-selection fallback on the same source file. Use `--writes-per-iteration` to repeat the same shard write within one timed run when you want to amplify per-row writer overhead. +For cloud runs, the summary also records `stage_duration_s`, the sum of stage +durations. That is often a better performance signal than wall time because it +excludes queueing noise from the cloud scheduler. + Artifacts are written under `benchmark/rerun/artifacts/` by default: - one per-case result JSON diff --git a/benchmark/rerun/compare_results.py b/benchmark/rerun/compare_results.py index 7c92e443..6daea535 100644 --- a/benchmark/rerun/compare_results.py +++ b/benchmark/rerun/compare_results.py @@ -94,6 +94,27 @@ def _stage_means(results: Sequence[Mapping[str, Any]]) -> dict[str, float | None return {stage: _mean(values) for stage, values in durations.items()} +def _stage_total(results: Sequence[Mapping[str, Any]]) -> float | None: + values = [] + for result in results: + total = result.get("stage_duration_s") + if isinstance(total, (int, float)): + values.append(float(total)) + continue + stages = result.get("stage_results") + if not isinstance(stages, list): + continue + durations = [ + float(stage["duration_s"]) + for stage in stages + if isinstance(stage, dict) + and isinstance(stage.get("duration_s"), (int, float)) + ] + if durations: + values.append(sum(durations)) + return _mean(values) + + def _delta( baseline: float | None, candidate: float | None, @@ -124,7 +145,12 @@ def _comparison( candidate_wall = _mean( result.get("cloud_wall_time_s") for result in candidate_completed ) + baseline_stage_total = _stage_total(baseline_completed) + candidate_stage_total = _stage_total(candidate_completed) wall_delta_s, wall_delta_pct = _delta(baseline_wall, candidate_wall) + stage_total_delta_s, stage_total_delta_pct = _delta( + baseline_stage_total, candidate_stage_total + ) baseline_stages = _stage_means(baseline_completed) candidate_stages = _stage_means(candidate_completed) stage_rows = [] @@ -164,6 +190,10 @@ def _comparison( "candidate_wall_time_s": candidate_wall, "delta_s": wall_delta_s, "delta_pct": wall_delta_pct, + "baseline_stage_total_s": baseline_stage_total, + "candidate_stage_total_s": candidate_stage_total, + "stage_total_delta_s": stage_total_delta_s, + "stage_total_delta_pct": stage_total_delta_pct, "stages": stage_rows, } ) @@ -218,6 +248,7 @@ def _print_human(comparison: Mapping[str, Any]) -> None: "candidate_s", "delta_s", "delta_pct", + "stage_total_s", ) ] for case in comparison["cases"]: @@ -232,9 +263,13 @@ def _print_human(comparison: Mapping[str, Any]) -> None: _format_number(case["candidate_wall_time_s"]), _format_number(case["delta_s"]), _format_number(case["delta_pct"], suffix="%"), + f"{_format_number(case['baseline_stage_total_s'])} -> " + f"{_format_number(case['candidate_stage_total_s'])}", ) ) _print_table(rows) + print() + print("stage_total = sum of stage durations; wall_time = cloud job elapsed time") for case in comparison["cases"]: stages = case["stages"] diff --git a/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py index 93ef658f..99a7d517 100644 --- a/benchmark/rerun/run_cloud_benchmark.py +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -66,6 +66,7 @@ class CaseResult: output_root: str cloud_wall_time_s: float | None queue_time_s: float | None + stage_duration_s: float | None stage_results: list[StageResult] output_file_count: int | None output_size_bytes: int | None @@ -402,6 +403,17 @@ def _stage_results(client: MacrodataClient, job: dict[str, Any]) -> list[StageRe return out +def _stage_duration_s(stage_results: Sequence[StageResult]) -> float | None: + durations = [ + stage.duration_s + for stage in stage_results + if isinstance(stage.duration_s, (int, float)) + ] + if not durations: + return None + return float(sum(durations)) + + def _optional_int(value: Any) -> int | None: return int(value) if isinstance(value, (int, float)) else None @@ -502,6 +514,7 @@ def _run_case( output_error: str | None = None if not args.skip_output_inspection: output_file_count, output_size_bytes, output_error = _inspect_output(output) + stage_results = _stage_results(client, job) return CaseResult( case=case, @@ -517,7 +530,8 @@ def _run_case( output_root=output, cloud_wall_time_s=_duration_s(job.get("startedAt"), job.get("endedAt")), queue_time_s=_duration_s(job.get("createdAt"), job.get("startedAt")), - stage_results=_stage_results(client, job), + stage_duration_s=_stage_duration_s(stage_results), + stage_results=stage_results, output_file_count=output_file_count, output_size_bytes=output_size_bytes, output_inspection_error=output_error, diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index f9217593..f379b1f6 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -105,6 +105,16 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: return count def _write_recording(self, recording: RerunRecording, relpath: str) -> None: + target = self.output.file(relpath) + if ( + self.write_footer + and recording.use_source_chunks + and recording.source_file is not None + ): + if _can_copy_source_rrd(recording): + recording.source_file.copy(target) + return + def write_local(path: Path | str) -> None: if ( self.write_footer @@ -127,7 +137,6 @@ def write_local(path: Path | str) -> None: write_local(local_path) return - target = self.output.file(relpath) with tempfile.TemporaryDirectory(prefix="refiner-rerun-write-") as tmpdir: local_path = Path(tmpdir) / os.path.basename(relpath) write_local(local_path) diff --git a/src/refiner/platform/manifest.py b/src/refiner/platform/manifest.py index ac9635a1..82e0f493 100644 --- a/src/refiner/platform/manifest.py +++ b/src/refiner/platform/manifest.py @@ -226,6 +226,22 @@ def _resolve_local_repo_git_sha() -> str | None: def refiner_ref_exists_on_remote(ref: str) -> bool: + try: + subprocess.run( + [ + "gh", + "api", + f"repos/macrodata-labs/refiner/commits/{ref}", + "--silent", + ], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + return True + except (FileNotFoundError, subprocess.CalledProcessError): + pass + request = urllib_request.Request( f"https://api.github.com/repos/macrodata-labs/refiner/commits/{ref}" ) diff --git a/tests/benchmark/test_compare_results.py b/tests/benchmark/test_compare_results.py new file mode 100644 index 00000000..8cc0a4ad --- /dev/null +++ b/tests/benchmark/test_compare_results.py @@ -0,0 +1,80 @@ +from __future__ import annotations + +import importlib.util +from pathlib import Path +from types import ModuleType + + +def _load_compare_results() -> ModuleType: + path = ( + Path(__file__).resolve().parents[2] + / "benchmark" + / "rerun" + / "compare_results.py" + ) + spec = importlib.util.spec_from_file_location("compare_results", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +compare_results = _load_compare_results() + + +def test_stage_total_is_derived_from_stage_durations() -> None: + summary = { + "results": [ + { + "case": "rrd-copy", + "status": "completed", + "cloud_wall_time_s": 66.07, + "stage_results": [ + {"name": "write_rerun_stage_0", "duration_s": 14.88}, + {"name": "write_rerun_stage_1", "duration_s": 3.65}, + ], + } + ] + } + + assert compare_results._stage_total(summary["results"]) == 18.53 + + +def test_comparison_reports_stage_total_delta() -> None: + baseline = { + "run_token": "baseline", + "git_ref": "base", + "results": [ + { + "case": "rrd-copy", + "status": "completed", + "cloud_wall_time_s": 61.82, + "stage_results": [ + {"name": "write_rerun_stage_0", "duration_s": 20.80}, + {"name": "write_rerun_stage_1", "duration_s": 2.32}, + ], + } + ], + } + candidate = { + "run_token": "candidate", + "git_ref": "cand", + "results": [ + { + "case": "rrd-copy", + "status": "completed", + "cloud_wall_time_s": 66.07, + "stage_results": [ + {"name": "write_rerun_stage_0", "duration_s": 14.88}, + {"name": "write_rerun_stage_1", "duration_s": 3.65}, + ], + } + ], + } + + comparison = compare_results._comparison(baseline, candidate) + case = comparison["cases"][0] + + assert case["baseline_stage_total_s"] == 23.12 + assert case["candidate_stage_total_s"] == 18.53 + assert case["stage_total_delta_s"] == -4.59 diff --git a/tests/platform/test_manifest.py b/tests/platform/test_manifest.py index ed456f1a..ceaf60b5 100644 --- a/tests/platform/test_manifest.py +++ b/tests/platform/test_manifest.py @@ -6,6 +6,7 @@ from email.message import Message from importlib import metadata as importlib_metadata from pathlib import Path +from typing import cast from urllib import error as urllib_error import pytest @@ -461,6 +462,10 @@ def test_build_run_manifest_environment_does_not_include_rundir_by_default( def test_refiner_ref_exists_on_remote_returns_true_on_success(monkeypatch) -> None: + def _raise_no_gh(*args, **kwargs): + raise FileNotFoundError + + monkeypatch.setattr("refiner.platform.manifest.subprocess.run", _raise_no_gh) monkeypatch.setattr( "refiner.platform.manifest.urllib_request.urlopen", lambda request: nullcontext(object()), @@ -470,6 +475,11 @@ def test_refiner_ref_exists_on_remote_returns_true_on_success(monkeypatch) -> No def test_refiner_ref_exists_on_remote_returns_false_on_404(monkeypatch) -> None: + def _raise_no_gh(*args, **kwargs): + raise FileNotFoundError + + monkeypatch.setattr("refiner.platform.manifest.subprocess.run", _raise_no_gh) + def _raise_404(request): raise urllib_error.HTTPError( request.full_url, @@ -487,6 +497,29 @@ def _raise_404(request): assert refiner_ref_exists_on_remote("abc123") is False +def test_refiner_ref_exists_on_remote_prefers_gh_api(monkeypatch) -> None: + calls: list[tuple[tuple[object, ...], dict[str, object]]] = [] + + def _fake_run(*args, **kwargs): + calls.append((args, kwargs)) + return None + + monkeypatch.setattr("refiner.platform.manifest.subprocess.run", _fake_run) + monkeypatch.setattr( + "refiner.platform.manifest.urllib_request.urlopen", + lambda request: pytest.fail("urllib fallback should not be used when gh works"), + ) + + assert refiner_ref_exists_on_remote("abc123") is True + assert calls + command = cast(list[str], calls[0][0][0]) + assert command[:3] == [ + "gh", + "api", + "repos/macrodata-labs/refiner/commits/abc123", + ] + + def test_manifest_prefers_macrodata_refiner_distribution(monkeypatch) -> None: def _version(name: str) -> str: if name == "macrodata-refiner": diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index fa16503a..898c5e7e 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -741,6 +741,50 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert copied["rerun"].tables["frame"].num_rows == 3 +def test_write_rerun_copies_source_file_directly_for_raw_copy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "tiny.rrd" + _tiny_rrd(source) + + source_iter = mdr.read_rerun( + str(source), + materialize_tables=False, + ).source.read() + block = cast(list[Row], next(source_iter)) + + output_fs = fsspec.filesystem("memory") + output = ("bucket/out-raw-copy-direct", output_fs) + + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.LocalRrd", + lambda *args, **kwargs: pytest.fail( + "raw source chunks should not stage a local RRD copy" + ), + ) + monkeypatch.setattr( + "refiner.pipeline.sinks.rerun.tempfile.TemporaryDirectory", + lambda *args, **kwargs: pytest.fail( + "raw source chunks should copy straight to the final target" + ), + ) + + sink = RerunSink(output) + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + with pytest.raises(StopIteration): + next(source_iter) + + written = [ + path + for path in output_fs.find("bucket/out-raw-copy-direct") + if path.endswith(".rrd") + ] + assert len(written) == 1 + assert output_fs.cat(written[0]) == source.read_bytes() + + def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From 9a8cc381e77a1edacc10984d809743c4035080bd Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:07:41 +0200 Subject: [PATCH 53/65] Prefer native remote RRD staging --- src/refiner/pipeline/_rerun_io.py | 7 ++++ src/refiner/pipeline/sinks/rerun.py | 8 ----- tests/readers/test_rerun_reader.py | 53 ++++++++++++----------------- 3 files changed, 28 insertions(+), 40 deletions(-) diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 5388c364..446864ae 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -26,6 +26,13 @@ def open(self) -> Path: self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") name = os.path.basename(self.source.path) or "recording.rrd" self.path = Path(self.tmpdir.name) / name + get_file = getattr(self.source.fs, "get_file", None) + if callable(get_file): + try: + get_file(self.source.path, str(self.path)) + return self.path + except Exception: + pass self.source.copy(str(self.path)) return self.path diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py index f379b1f6..ef49270e 100644 --- a/src/refiner/pipeline/sinks/rerun.py +++ b/src/refiner/pipeline/sinks/rerun.py @@ -106,14 +106,6 @@ def write_shard_block(self, shard_id: str, block: Block) -> int: def _write_recording(self, recording: RerunRecording, relpath: str) -> None: target = self.output.file(relpath) - if ( - self.write_footer - and recording.use_source_chunks - and recording.source_file is not None - ): - if _can_copy_source_rrd(recording): - recording.source_file.copy(target) - return def write_local(path: Path | str) -> None: if ( diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index 898c5e7e..ae38d711 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -10,7 +10,9 @@ import pytest import refiner as mdr +from refiner.io.datafile import DataFile from refiner.pipeline import Row +from refiner.pipeline._rerun_io import LocalRrd from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table @@ -741,48 +743,35 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert copied["rerun"].tables["frame"].num_rows == 3 -def test_write_rerun_copies_source_file_directly_for_raw_copy( +def test_local_rrd_prefers_get_file_for_remote_sources( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - source = tmp_path / "tiny.rrd" - _tiny_rrd(source) + source = tmp_path / "remote.rrd" + source.write_bytes(b"rrd") - source_iter = mdr.read_rerun( - str(source), - materialize_tables=False, - ).source.read() - block = cast(list[Row], next(source_iter)) + remote_fs = fsspec.filesystem("memory") + remote_path = "/refiner-rerun-test/remote.rrd" + remote_fs.pipe_file(remote_path, source.read_bytes()) - output_fs = fsspec.filesystem("memory") - output = ("bucket/out-raw-copy-direct", output_fs) + calls: list[tuple[str, str, dict[str, Any]]] = [] + original_get_file = remote_fs.get_file + def fake_get_file(src: str, dst: str, **kwargs: Any): + calls.append((src, dst, dict(kwargs))) + return original_get_file(src, dst, **kwargs) + + monkeypatch.setattr(remote_fs, "get_file", fake_get_file) monkeypatch.setattr( - "refiner.pipeline.sinks.rerun.LocalRrd", - lambda *args, **kwargs: pytest.fail( - "raw source chunks should not stage a local RRD copy" - ), - ) - monkeypatch.setattr( - "refiner.pipeline.sinks.rerun.tempfile.TemporaryDirectory", - lambda *args, **kwargs: pytest.fail( - "raw source chunks should copy straight to the final target" - ), + "refiner.io.datafile.DataFile.copy", + lambda *args, **kwargs: pytest.fail("get_file should be preferred"), ) - sink = RerunSink(output) - sink.write_shard_block("shard-a", block) - sink.on_shard_complete("shard-a") - with pytest.raises(StopIteration): - next(source_iter) + local_rrd = LocalRrd(DataFile.resolve((remote_path, remote_fs))) + path = local_rrd.open() - written = [ - path - for path in output_fs.find("bucket/out-raw-copy-direct") - if path.endswith(".rrd") - ] - assert len(written) == 1 - assert output_fs.cat(written[0]) == source.read_bytes() + assert path.exists() + assert calls == [(remote_path, str(path), {})] def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( From abbf4a76413a92ef75bcf79198ff829ddd3360c8 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:18:41 +0200 Subject: [PATCH 54/65] Revert native remote RRD staging --- src/refiner/pipeline/_rerun_io.py | 7 ------- tests/readers/test_rerun_reader.py | 33 ------------------------------ 2 files changed, 40 deletions(-) diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 446864ae..5388c364 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -26,13 +26,6 @@ def open(self) -> Path: self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") name = os.path.basename(self.source.path) or "recording.rrd" self.path = Path(self.tmpdir.name) / name - get_file = getattr(self.source.fs, "get_file", None) - if callable(get_file): - try: - get_file(self.source.path, str(self.path)) - return self.path - except Exception: - pass self.source.copy(str(self.path)) return self.path diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index ae38d711..fa16503a 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -10,9 +10,7 @@ import pytest import refiner as mdr -from refiner.io.datafile import DataFile from refiner.pipeline import Row -from refiner.pipeline._rerun_io import LocalRrd from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table @@ -743,37 +741,6 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert copied["rerun"].tables["frame"].num_rows == 3 -def test_local_rrd_prefers_get_file_for_remote_sources( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - source = tmp_path / "remote.rrd" - source.write_bytes(b"rrd") - - remote_fs = fsspec.filesystem("memory") - remote_path = "/refiner-rerun-test/remote.rrd" - remote_fs.pipe_file(remote_path, source.read_bytes()) - - calls: list[tuple[str, str, dict[str, Any]]] = [] - original_get_file = remote_fs.get_file - - def fake_get_file(src: str, dst: str, **kwargs: Any): - calls.append((src, dst, dict(kwargs))) - return original_get_file(src, dst, **kwargs) - - monkeypatch.setattr(remote_fs, "get_file", fake_get_file) - monkeypatch.setattr( - "refiner.io.datafile.DataFile.copy", - lambda *args, **kwargs: pytest.fail("get_file should be preferred"), - ) - - local_rrd = LocalRrd(DataFile.resolve((remote_path, remote_fs))) - path = local_rrd.open() - - assert path.exists() - assert calls == [(remote_path, str(path), {})] - - def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From e15f2b7510e3ae450b54d2543fe5c4c3d16c1b66 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:19:25 +0200 Subject: [PATCH 55/65] Increase RRD staging buffer size --- src/refiner/pipeline/_rerun_io.py | 2 +- tests/readers/test_rerun_reader.py | 29 +++++++++++++++++++++++++++++ 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 5388c364..84ba426d 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -26,7 +26,7 @@ def open(self) -> Path: self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") name = os.path.basename(self.source.path) or "recording.rrd" self.path = Path(self.tmpdir.name) / name - self.source.copy(str(self.path)) + self.source.copy(str(self.path), buffer_size=8 * 1024 * 1024) return self.path def close(self) -> None: diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index fa16503a..b40e9d4a 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -10,7 +10,9 @@ import pytest import refiner as mdr +from refiner.io.datafile import DataFile from refiner.pipeline import Row +from refiner.pipeline._rerun_io import LocalRrd from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table @@ -741,6 +743,33 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert copied["rerun"].tables["frame"].num_rows == 3 +def test_local_rrd_uses_larger_buffer_for_remote_sources( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + source = tmp_path / "remote.rrd" + source.write_bytes(b"rrd") + + remote_fs = fsspec.filesystem("memory") + remote_path = "/refiner-rerun-test/remote.rrd" + remote_fs.pipe_file(remote_path, source.read_bytes()) + + buffer_sizes: list[int] = [] + original_copy = DataFile.copy + + def fake_copy(self, dest, *, buffer_size: int = 2 * 1024 * 1024): + buffer_sizes.append(buffer_size) + return original_copy(self, dest, buffer_size=buffer_size) + + monkeypatch.setattr(DataFile, "copy", fake_copy) + + local_rrd = LocalRrd(DataFile.resolve((remote_path, remote_fs))) + path = local_rrd.open() + + assert path.exists() + assert buffer_sizes == [8 * 1024 * 1024] + + def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From 45a13fb0fca1bfca7f0aa1c17927b2cc4d0df10e Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:28:43 +0200 Subject: [PATCH 56/65] Raise RRD reader fanout --- src/refiner/pipeline/_rerun_io.py | 2 +- src/refiner/pipeline/sources/readers/rerun.py | 4 +-- tests/readers/test_rerun_reader.py | 29 ------------------- 3 files changed, 3 insertions(+), 32 deletions(-) diff --git a/src/refiner/pipeline/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py index 84ba426d..5388c364 100644 --- a/src/refiner/pipeline/_rerun_io.py +++ b/src/refiner/pipeline/_rerun_io.py @@ -26,7 +26,7 @@ def open(self) -> Path: self.tmpdir = tempfile.TemporaryDirectory(prefix="refiner-rerun-") name = os.path.basename(self.source.path) or "recording.rrd" self.path = Path(self.tmpdir.name) / name - self.source.copy(str(self.path), buffer_size=8 * 1024 * 1024) + self.source.copy(str(self.path)) return self.path def close(self) -> None: diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 673de053..bde0bf32 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -678,7 +678,7 @@ def _open_local_sources( ] local_sources = [local_source for _source, local_source in local_files] - max_workers = min(8, len(local_files)) + max_workers = min(12, len(local_files)) try: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: local_paths = list(pool.map(_open_local_source, local_sources)) @@ -708,7 +708,7 @@ def _scan_recording_entries( ] local_paths = [local_path for _source, local_path, _local_source in local_files] - max_workers = min(8, len(local_files)) + max_workers = min(12, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: store_entries = list(pool.map(_recording_entries, local_paths)) return [ diff --git a/tests/readers/test_rerun_reader.py b/tests/readers/test_rerun_reader.py index b40e9d4a..fa16503a 100644 --- a/tests/readers/test_rerun_reader.py +++ b/tests/readers/test_rerun_reader.py @@ -10,9 +10,7 @@ import pytest import refiner as mdr -from refiner.io.datafile import DataFile from refiner.pipeline import Row -from refiner.pipeline._rerun_io import LocalRrd from refiner.pipeline.data.row import DictRow from refiner.pipeline.sinks.rerun import RerunSink from refiner.pipeline.sinks.rerun import _sendable_dynamic_table, _sendable_static_table @@ -743,33 +741,6 @@ def test_write_rerun_reuses_reader_staged_remote_source( assert copied["rerun"].tables["frame"].num_rows == 3 -def test_local_rrd_uses_larger_buffer_for_remote_sources( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - source = tmp_path / "remote.rrd" - source.write_bytes(b"rrd") - - remote_fs = fsspec.filesystem("memory") - remote_path = "/refiner-rerun-test/remote.rrd" - remote_fs.pipe_file(remote_path, source.read_bytes()) - - buffer_sizes: list[int] = [] - original_copy = DataFile.copy - - def fake_copy(self, dest, *, buffer_size: int = 2 * 1024 * 1024): - buffer_sizes.append(buffer_size) - return original_copy(self, dest, buffer_size=buffer_size) - - monkeypatch.setattr(DataFile, "copy", fake_copy) - - local_rrd = LocalRrd(DataFile.resolve((remote_path, remote_fs))) - path = local_rrd.open() - - assert path.exists() - assert buffer_sizes == [8 * 1024 * 1024] - - def test_write_rerun_does_not_direct_copy_when_source_may_have_multiple_recordings( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, From 6c71e6ea07bfecc2de7c80bbf8a376b854e7d360 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:36:51 +0200 Subject: [PATCH 57/65] Back off RRD metadata scan fanout --- src/refiner/pipeline/sources/readers/rerun.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index bde0bf32..20045e6a 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -678,7 +678,7 @@ def _open_local_sources( ] local_sources = [local_source for _source, local_source in local_files] - max_workers = min(12, len(local_files)) + max_workers = min(8, len(local_files)) try: with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: local_paths = list(pool.map(_open_local_source, local_sources)) From f9586eaefe6d5879fe7242eb651e14bb0698a745 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:45:49 +0200 Subject: [PATCH 58/65] Revert RRD reader fanout --- src/refiner/pipeline/sources/readers/rerun.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/refiner/pipeline/sources/readers/rerun.py b/src/refiner/pipeline/sources/readers/rerun.py index 20045e6a..673de053 100644 --- a/src/refiner/pipeline/sources/readers/rerun.py +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -708,7 +708,7 @@ def _scan_recording_entries( ] local_paths = [local_path for _source, local_path, _local_source in local_files] - max_workers = min(12, len(local_files)) + max_workers = min(8, len(local_files)) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as pool: store_entries = list(pool.map(_recording_entries, local_paths)) return [ From 170c51e2e4e5298ee0103dac266fe6da1694599a Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:48:01 +0200 Subject: [PATCH 59/65] Optimize default RRD cleanup listing --- src/refiner/pipeline/sinks/reducer/file.py | 25 ++++++++++++ tests/pipeline/test_sinks.py | 47 ++++++++++++++++++++++ 2 files changed, 72 insertions(+) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index f32b079f..310cbf16 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -126,6 +126,31 @@ def _run_cleanup(self) -> None: listing_prefix = ( "" if "/" not in literal_prefix else literal_prefix.rsplit("/", 1)[0] ) + if ( + self.assets_subdir is None + and listing_prefix == "" + and len(self._output_path_patterns) == 2 + ): + try: + root_entries = self.output.ls(listing_prefix, detail=False) + except (FileNotFoundError, NotADirectoryError): + root_entries = [] + paths_to_delete: set[str] = set() + for rel_path in root_entries: + match = self._output_path_patterns[0].fullmatch(rel_path) + if match is None: + continue + if ( + match.group("shard_id"), + match.group("worker_id"), + ) not in keep_pairs: + paths_to_delete.add(rel_path) + for path in sorted(paths_to_delete): + try: + self.output.rm(path, recursive=True) + except FileNotFoundError: + continue + return paths = [listing_prefix] prefix_parts = [part for part in listing_prefix.split("/") if part] for pattern in self._output_path_patterns[len(prefix_parts) :]: diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 07ef0d55..06b487ba 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -964,6 +964,53 @@ def test_file_cleanup_reducer_removes_non_finalized_directories(tmp_path) -> Non assert not loser_dir.exists() +def test_file_cleanup_reducer_lists_root_once_for_default_rrd_layout( + tmp_path, + monkeypatch, +) -> None: + output_dir = tmp_path / "rrd-cleanup" + shard_id = "0123456789ab" + winner_worker_id = "worker-2" + loser_worker_id = "worker-1" + winner_dir = output_dir / f"{shard_id}__w{worker_token_for(winner_worker_id)}" + loser_dir = output_dir / f"{shard_id}__w{worker_token_for(loser_worker_id)}" + winner_dir.mkdir(parents=True) + loser_dir.mkdir(parents=True) + (winner_dir / "0.rrd").write_bytes(b"keep") + (loser_dir / "0.rrd").write_bytes(b"drop") + + reducer = FileCleanupReducerSink( + output_dir, + filename_template="{shard_id}__w{worker_id}/{row_index}.rrd", + reducer_name="cleanup_rrd", + ) + ls_calls: list[str] = [] + original_ls = reducer.output.ls + + def fake_ls(path, detail=False): + ls_calls.append(path) + return original_ls(path, detail=detail) + + monkeypatch.setattr(reducer.output, "ls", fake_ls) + with set_active_run_context( + job_id="job", + stage_index=1, + worker_id="reducer", + worker_name=None, + runtime_lifecycle=cast( + RuntimeLifecycle, + _FinalizedWorkersRuntime( + [FinalizedShardWorker(shard_id=shard_id, worker_id=winner_worker_id)] + ), + ), + ): + reducer.write_block([DictRow({"task_rank": 0}, shard_id="reduce")]) + + assert len(ls_calls) == 1 + assert winner_dir.exists() + assert not loser_dir.exists() + + def test_file_cleanup_reducer_removes_non_finalized_nested_directories( tmp_path, ) -> None: From 673fdc144f7e18978aa64c49f6e835ebae4d50b3 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:54:33 +0200 Subject: [PATCH 60/65] Speed up default RRD cleanup matching --- src/refiner/pipeline/sinks/reducer/file.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index 310cbf16..6dcbb14d 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -137,13 +137,11 @@ def _run_cleanup(self) -> None: root_entries = [] paths_to_delete: set[str] = set() for rel_path in root_entries: - match = self._output_path_patterns[0].fullmatch(rel_path) - if match is None: + name = rel_path.rstrip("/") + shard_id, sep, worker_id = name.partition("__w") + if not sep or len(shard_id) != 12 or len(worker_id) != 12: continue - if ( - match.group("shard_id"), - match.group("worker_id"), - ) not in keep_pairs: + if (shard_id, worker_id) not in keep_pairs: paths_to_delete.add(rel_path) for path in sorted(paths_to_delete): try: From 095c66c5e088f96937def569c646df81e66f866a Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:55:56 +0200 Subject: [PATCH 61/65] Tighten default RRD cleanup parsing --- src/refiner/pipeline/sinks/reducer/file.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index 6dcbb14d..02e029cf 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -137,10 +137,10 @@ def _run_cleanup(self) -> None: root_entries = [] paths_to_delete: set[str] = set() for rel_path in root_entries: - name = rel_path.rstrip("/") - shard_id, sep, worker_id = name.partition("__w") - if not sep or len(shard_id) != 12 or len(worker_id) != 12: + if len(rel_path) != 27 or rel_path[12:15] != "__w": continue + shard_id = rel_path[:12] + worker_id = rel_path[15:] if (shard_id, worker_id) not in keep_pairs: paths_to_delete.add(rel_path) for path in sorted(paths_to_delete): From f04902720a47496a52c36602b92400a420c2572e Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:57:50 +0200 Subject: [PATCH 62/65] Add cleanup matcher benchmark --- benchmark/rerun/README.md | 8 ++ benchmark/rerun/run_cleanup_benchmark.py | 141 +++++++++++++++++++++ src/refiner/pipeline/sinks/reducer/file.py | 24 ++-- tests/pipeline/test_sinks.py | 14 ++ 4 files changed, 179 insertions(+), 8 deletions(-) create mode 100644 benchmark/rerun/run_cleanup_benchmark.py diff --git a/benchmark/rerun/README.md b/benchmark/rerun/README.md index 595e86c1..63e6a8b5 100644 --- a/benchmark/rerun/README.md +++ b/benchmark/rerun/README.md @@ -11,6 +11,8 @@ This folder contains cloud benchmark harnesses for Rerun RRD workloads. profile into the Macrodata workspace secret environment used by cloud jobs. - `run_local_benchmark.py`: runs a local single-recording RRD copy benchmark that compares the direct byte-copy path with the chunk-selection fallback. +- `run_cleanup_benchmark.py`: runs a local benchmark for the default-root RRD + cleanup matcher used by `FileCleanupReducerSink`. The default inputs are the ten base RRD files from: @@ -76,6 +78,12 @@ For a local smoke benchmark that does not require cloud credentials: uv run python benchmark/rerun/run_local_benchmark.py ``` +For the reducer cleanup matcher benchmark: + +```bash +uv run python benchmark/rerun/run_cleanup_benchmark.py +``` + The local benchmark generates a synthetic single-recording RRD, then measures the direct-copy branch against the chunk-selection fallback on the same source file. Use `--writes-per-iteration` to repeat the same shard write within one diff --git a/benchmark/rerun/run_cleanup_benchmark.py b/benchmark/rerun/run_cleanup_benchmark.py new file mode 100644 index 00000000..045bd38d --- /dev/null +++ b/benchmark/rerun/run_cleanup_benchmark.py @@ -0,0 +1,141 @@ +from __future__ import annotations + +import argparse +import json +import re +from dataclasses import asdict, dataclass +from datetime import datetime, timezone +from pathlib import Path +from time import perf_counter_ns +from typing import Callable + +from refiner.pipeline.sinks.reducer.file import _cleanup_default_root_entries + +DEFAULT_ARTIFACTS_DIR = Path(__file__).resolve().parent / "artifacts" +_REGEX_PATTERN = re.compile( + r"^(?P[0-9a-f]{12})__w(?P[0-9a-f]{12})$" +) + + +@dataclass(slots=True) +class CaseResult: + mode: str + wall_time_ns: int + entries: int + deleted_entries: int + + @property + def wall_time_s(self) -> float: + return self.wall_time_ns / 1_000_000_000 + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Benchmark the default-root RRD cleanup matcher." + ) + parser.add_argument("--iterations", type=int, default=200) + parser.add_argument("--entries", type=int, default=10000) + parser.add_argument("--artifacts-dir", type=Path, default=DEFAULT_ARTIFACTS_DIR) + parser.add_argument("--run-token") + return parser.parse_args() + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat() + + +def _generate_entries(entries: int) -> tuple[list[str], set[tuple[str, str]]]: + root_entries = [f"{index:012x}__w{(index + 1):012x}" for index in range(entries)] + keep_pairs = { + (f"{index:012x}", f"{(index + 1):012x}") for index in range(0, entries, 2) + } + return root_entries, keep_pairs + + +def _regex_cleanup( + root_entries: list[str], + keep_pairs: set[tuple[str, str]], +) -> set[str]: + paths_to_delete: set[str] = set() + for rel_path in root_entries: + match = _REGEX_PATTERN.fullmatch(rel_path) + if match is None: + continue + if (match.group("shard_id"), match.group("worker_id")) not in keep_pairs: + paths_to_delete.add(rel_path) + return paths_to_delete + + +def _benchmark( + *, + mode: str, + fn: Callable[[list[str], set[tuple[str, str]]], set[str]], + root_entries: list[str], + keep_pairs: set[tuple[str, str]], + iterations: int, +) -> CaseResult: + start = perf_counter_ns() + deleted_entries = 0 + for _ in range(iterations): + deleted_entries = len(fn(root_entries, keep_pairs)) + return CaseResult( + mode=mode, + wall_time_ns=perf_counter_ns() - start, + entries=len(root_entries), + deleted_entries=deleted_entries, + ) + + +def main() -> int: + args = _parse_args() + if args.iterations < 1: + raise ValueError("--iterations must be >= 1") + if args.entries < 1: + raise ValueError("--entries must be >= 1") + + run_token = ( + args.run_token + or f"cleanup-{datetime.now(timezone.utc).strftime('%Y%m%dT%H%M%SZ')}" + ) + artifacts_dir = args.artifacts_dir / run_token + artifacts_dir.mkdir(parents=True, exist_ok=True) + + root_entries, keep_pairs = _generate_entries(args.entries) + results = [ + _benchmark( + mode="regex", + fn=_regex_cleanup, + root_entries=root_entries, + keep_pairs=keep_pairs, + iterations=args.iterations, + ), + _benchmark( + mode="fixed-slice", + fn=_cleanup_default_root_entries, + root_entries=root_entries, + keep_pairs=keep_pairs, + iterations=args.iterations, + ), + ] + + summary = { + "run_token": run_token, + "started_at_utc": _utc_now(), + "iterations": args.iterations, + "entries": args.entries, + "results": [asdict(result) for result in results], + } + summary_path = artifacts_dir / "summary.json" + summary_path.write_text( + json.dumps(summary, indent=2, sort_keys=True), encoding="utf-8" + ) + print(f"Summary written to {summary_path}") + for result in results: + print( + f"{result.mode}: {result.wall_time_s:.6f}s deleted={result.deleted_entries}" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index 02e029cf..a0f8172b 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -135,14 +135,7 @@ def _run_cleanup(self) -> None: root_entries = self.output.ls(listing_prefix, detail=False) except (FileNotFoundError, NotADirectoryError): root_entries = [] - paths_to_delete: set[str] = set() - for rel_path in root_entries: - if len(rel_path) != 27 or rel_path[12:15] != "__w": - continue - shard_id = rel_path[:12] - worker_id = rel_path[15:] - if (shard_id, worker_id) not in keep_pairs: - paths_to_delete.add(rel_path) + paths_to_delete = _cleanup_default_root_entries(root_entries, keep_pairs) for path in sorted(paths_to_delete): try: self.output.rm(path, recursive=True) @@ -201,3 +194,18 @@ def _run_cleanup(self) -> None: __all__ = ["FileCleanupReducerSink"] + + +def _cleanup_default_root_entries( + root_entries: list[str], + keep_pairs: set[tuple[str, str]], +) -> set[str]: + paths_to_delete: set[str] = set() + for rel_path in root_entries: + if len(rel_path) != 27 or rel_path[12:15] != "__w": + continue + shard_id = rel_path[:12] + worker_id = rel_path[15:] + if (shard_id, worker_id) not in keep_pairs: + paths_to_delete.add(rel_path) + return paths_to_delete diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 06b487ba..f3cdcbe6 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -16,6 +16,7 @@ from refiner.pipeline.sinks import JsonlSink from refiner.pipeline.sinks.parquet import ParquetSink from refiner.pipeline.sinks.reducer.file import FileCleanupReducerSink +from refiner.pipeline.sinks.reducer.file import _cleanup_default_root_entries from refiner.worker.context import set_active_run_context from refiner.worker.lifecycle import FinalizedShardWorker, RuntimeLifecycle from refiner.worker.context import worker_token_for @@ -1011,6 +1012,19 @@ def fake_ls(path, detail=False): assert not loser_dir.exists() +def test_cleanup_default_root_entries_returns_only_losers() -> None: + root_entries = [ + "0123456789ab__w111111111111", + "0123456789ab__w222222222222", + "not-a-match", + ] + keep_pairs = {("0123456789ab", "111111111111")} + + assert _cleanup_default_root_entries(root_entries, keep_pairs) == { + "0123456789ab__w222222222222" + } + + def test_file_cleanup_reducer_removes_non_finalized_nested_directories( tmp_path, ) -> None: From 7fe8beef63c2368b831d044b3db8c10a8b0abb83 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 20:58:55 +0200 Subject: [PATCH 63/65] Speed up cleanup key lookup --- src/refiner/pipeline/sinks/reducer/file.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index a0f8172b..bf4b1163 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -200,12 +200,11 @@ def _cleanup_default_root_entries( root_entries: list[str], keep_pairs: set[tuple[str, str]], ) -> set[str]: + keep_keys = {f"{shard_id}__w{worker_id}" for shard_id, worker_id in keep_pairs} paths_to_delete: set[str] = set() for rel_path in root_entries: if len(rel_path) != 27 or rel_path[12:15] != "__w": continue - shard_id = rel_path[:12] - worker_id = rel_path[15:] - if (shard_id, worker_id) not in keep_pairs: + if rel_path not in keep_keys: paths_to_delete.add(rel_path) return paths_to_delete From 83ff8f23282830c5e865df78e8398d8399174535 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 21:00:18 +0200 Subject: [PATCH 64/65] Precompute cleanup matcher keys --- benchmark/rerun/run_cleanup_benchmark.py | 15 ++++++++------- src/refiner/pipeline/sinks/reducer/file.py | 6 +++--- tests/pipeline/test_sinks.py | 4 ++-- 3 files changed, 13 insertions(+), 12 deletions(-) diff --git a/benchmark/rerun/run_cleanup_benchmark.py b/benchmark/rerun/run_cleanup_benchmark.py index 045bd38d..a8e3ebe1 100644 --- a/benchmark/rerun/run_cleanup_benchmark.py +++ b/benchmark/rerun/run_cleanup_benchmark.py @@ -54,14 +54,14 @@ def _generate_entries(entries: int) -> tuple[list[str], set[tuple[str, str]]]: def _regex_cleanup( root_entries: list[str], - keep_pairs: set[tuple[str, str]], + keep_keys: set[str], ) -> set[str]: paths_to_delete: set[str] = set() for rel_path in root_entries: match = _REGEX_PATTERN.fullmatch(rel_path) if match is None: continue - if (match.group("shard_id"), match.group("worker_id")) not in keep_pairs: + if f"{match.group('shard_id')}__w{match.group('worker_id')}" not in keep_keys: paths_to_delete.add(rel_path) return paths_to_delete @@ -69,15 +69,15 @@ def _regex_cleanup( def _benchmark( *, mode: str, - fn: Callable[[list[str], set[tuple[str, str]]], set[str]], + fn: Callable[[list[str], set[str]], set[str]], root_entries: list[str], - keep_pairs: set[tuple[str, str]], + keep_keys: set[str], iterations: int, ) -> CaseResult: start = perf_counter_ns() deleted_entries = 0 for _ in range(iterations): - deleted_entries = len(fn(root_entries, keep_pairs)) + deleted_entries = len(fn(root_entries, keep_keys)) return CaseResult( mode=mode, wall_time_ns=perf_counter_ns() - start, @@ -101,19 +101,20 @@ def main() -> int: artifacts_dir.mkdir(parents=True, exist_ok=True) root_entries, keep_pairs = _generate_entries(args.entries) + keep_keys = {f"{shard_id}__w{worker_id}" for shard_id, worker_id in keep_pairs} results = [ _benchmark( mode="regex", fn=_regex_cleanup, root_entries=root_entries, - keep_pairs=keep_pairs, + keep_keys=keep_keys, iterations=args.iterations, ), _benchmark( mode="fixed-slice", fn=_cleanup_default_root_entries, root_entries=root_entries, - keep_pairs=keep_pairs, + keep_keys=keep_keys, iterations=args.iterations, ), ] diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index bf4b1163..78d5e353 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -115,6 +115,7 @@ def _run_cleanup(self) -> None: (row.shard_id, row.worker_token) for row in get_finalized_workers(stage_index=stage_index - 1) } + keep_keys = {f"{shard_id}__w{worker_id}" for shard_id, worker_id in keep_pairs} literal_prefix = "" for literal_text, field_name, _format_spec, _conversion in Formatter().parse( @@ -135,7 +136,7 @@ def _run_cleanup(self) -> None: root_entries = self.output.ls(listing_prefix, detail=False) except (FileNotFoundError, NotADirectoryError): root_entries = [] - paths_to_delete = _cleanup_default_root_entries(root_entries, keep_pairs) + paths_to_delete = _cleanup_default_root_entries(root_entries, keep_keys) for path in sorted(paths_to_delete): try: self.output.rm(path, recursive=True) @@ -198,9 +199,8 @@ def _run_cleanup(self) -> None: def _cleanup_default_root_entries( root_entries: list[str], - keep_pairs: set[tuple[str, str]], + keep_keys: set[str], ) -> set[str]: - keep_keys = {f"{shard_id}__w{worker_id}" for shard_id, worker_id in keep_pairs} paths_to_delete: set[str] = set() for rel_path in root_entries: if len(rel_path) != 27 or rel_path[12:15] != "__w": diff --git a/tests/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index f3cdcbe6..d2a11b80 100644 --- a/tests/pipeline/test_sinks.py +++ b/tests/pipeline/test_sinks.py @@ -1018,9 +1018,9 @@ def test_cleanup_default_root_entries_returns_only_losers() -> None: "0123456789ab__w222222222222", "not-a-match", ] - keep_pairs = {("0123456789ab", "111111111111")} + keep_keys = {"0123456789ab__w111111111111"} - assert _cleanup_default_root_entries(root_entries, keep_pairs) == { + assert _cleanup_default_root_entries(root_entries, keep_keys) == { "0123456789ab__w222222222222" } From 26b0a2b5d9044ec93a6271a92a1c5f16a58316b5 Mon Sep 17 00:00:00 2001 From: guipenedo Date: Mon, 15 Jun 2026 21:38:15 +0200 Subject: [PATCH 65/65] Restore cleanup best recipe --- src/refiner/pipeline/sinks/reducer/file.py | 58 ++++++++++++---------- src/refiner/worker/lifecycle.py | 9 ++-- tests/worker/test_runner.py | 7 +++ 3 files changed, 44 insertions(+), 30 deletions(-) diff --git a/src/refiner/pipeline/sinks/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index 78d5e353..7d0f3387 100644 --- a/src/refiner/pipeline/sinks/reducer/file.py +++ b/src/refiner/pipeline/sinks/reducer/file.py @@ -77,6 +77,21 @@ def __init__( self.reducer_name = reducer_name self.assets_subdir = assets_subdir self._output_path_patterns = _compile_output_path_patterns(filename_template) + literal_prefix = "" + for literal_text, field_name, _format_spec, _conversion in Formatter().parse( + self.filename_template + ): + literal_prefix += literal_text + if field_name is not None: + break + self._listing_prefix = ( + "" if "/" not in literal_prefix else literal_prefix.rsplit("/", 1)[0] + ) + self._cleanup_uses_default_root = ( + self.assets_subdir is None + and self._listing_prefix == "" + and len(self._output_path_patterns) == 2 + ) self._cleanup_ran = False def write_shard_block(self, shard_id, block) -> None: @@ -111,40 +126,31 @@ def _run_cleanup(self) -> None: f"{self.reducer_name} requires an active reducer stage with a prior writer stage" ) - keep_pairs = { - (row.shard_id, row.worker_token) - for row in get_finalized_workers(stage_index=stage_index - 1) - } - keep_keys = {f"{shard_id}__w{worker_id}" for shard_id, worker_id in keep_pairs} + finalized_workers = get_finalized_workers(stage_index=stage_index - 1) - literal_prefix = "" - for literal_text, field_name, _format_spec, _conversion in Formatter().parse( - self.filename_template - ): - literal_prefix += literal_text - if field_name is not None: - break - listing_prefix = ( - "" if "/" not in literal_prefix else literal_prefix.rsplit("/", 1)[0] - ) - if ( - self.assets_subdir is None - and listing_prefix == "" - and len(self._output_path_patterns) == 2 - ): + if self._cleanup_uses_default_root: + keep_keys = { + f"{row.shard_id}__w{row.worker_token}" for row in finalized_workers + } + rm = self.output.rm + keep_key_contains = keep_keys.__contains__ try: - root_entries = self.output.ls(listing_prefix, detail=False) + root_entries = self.output.ls(self._listing_prefix, detail=False) except (FileNotFoundError, NotADirectoryError): root_entries = [] - paths_to_delete = _cleanup_default_root_entries(root_entries, keep_keys) - for path in sorted(paths_to_delete): + for rel_path in root_entries: + if len(rel_path) != 27 or rel_path[12:15] != "__w": + continue + if keep_key_contains(rel_path): + continue try: - self.output.rm(path, recursive=True) + rm(rel_path, recursive=True) except FileNotFoundError: continue return - paths = [listing_prefix] - prefix_parts = [part for part in listing_prefix.split("/") if part] + keep_pairs = {(row.shard_id, row.worker_token) for row in finalized_workers} + paths = [self._listing_prefix] + prefix_parts = [part for part in self._listing_prefix.split("/") if part] for pattern in self._output_path_patterns[len(prefix_parts) :]: next_paths: list[str] = [] for path in paths: diff --git a/src/refiner/worker/lifecycle.py b/src/refiner/worker/lifecycle.py index 2d081876..8f41b6ab 100644 --- a/src/refiner/worker/lifecycle.py +++ b/src/refiner/worker/lifecycle.py @@ -10,14 +10,15 @@ from refiner.worker.context import worker_token_for -class FinalizedShardWorker(msgspec.Struct, frozen=True): +class FinalizedShardWorker(msgspec.Struct): shard_id: str worker_id: str global_ordinal: int | None = None + worker_token: str = "" - @property - def worker_token(self) -> str: - return worker_token_for(self.worker_id) + def __post_init__(self) -> None: + if not self.worker_token: + self.worker_token = worker_token_for(self.worker_id) class RuntimeLifecycle(Protocol): diff --git a/tests/worker/test_runner.py b/tests/worker/test_runner.py index 4ea541c3..2f82032f 100644 --- a/tests/worker/test_runner.py +++ b/tests/worker/test_runner.py @@ -20,6 +20,7 @@ from refiner.pipeline.data.row import DictRow, Row from refiner.worker.metrics.api import log_gauge from refiner.worker.lifecycle import FinalizedShardWorker, sort_finalized_workers +from refiner.worker.context import worker_token_for class _FakeReader(BaseReader): @@ -85,6 +86,12 @@ def test_sort_finalized_workers_uses_legacy_order_when_any_ordinal_is_missing() ] +def test_finalized_worker_caches_worker_token() -> None: + row = FinalizedShardWorker("shard-a", "worker-a") + + assert row.worker_token == worker_token_for("worker-a") + + class _NoopTelemetryEmitter: def emit_user_counter(self, **kwargs) -> None: del kwargs