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 new file mode 100644 index 00000000..63e6a8b5 --- /dev/null +++ b/benchmark/rerun/README.md @@ -0,0 +1,118 @@ +# 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. +- `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. +- `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: + +```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", 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 +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`. + 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: + +```bash +uv run python benchmark/rerun/refresh_aws_secrets.py \ + --aws-profile 210049840512_Researcher \ + --secret-env 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. +- `--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. + +For a local smoke benchmark that does not require cloud credentials: + +```bash +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 +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 +- 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: + +```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. +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 new file mode 100644 index 00000000..6daea535 --- /dev/null +++ b/benchmark/rerun/compare_results.py @@ -0,0 +1,324 @@ +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 _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: + 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 _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, +) -> 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 + ) + 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 = [] + 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_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, + "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, + } + ) + 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 _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): + 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", + "shards", + "baseline_s", + "candidate_s", + "delta_s", + "delta_pct", + "stage_total_s", + ) + ] + for case in comparison["cases"]: + rows.append( + ( + 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"]), + _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"] + 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) + + 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() + 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()) diff --git a/benchmark/rerun/refresh_aws_secrets.py b/benchmark/rerun/refresh_aws_secrets.py new file mode 100644 index 00000000..f8ecd56a --- /dev/null +++ b/benchmark/rerun/refresh_aws_secrets.py @@ -0,0 +1,140 @@ +from __future__ import annotations + +import argparse +import json +import os +import subprocess + + +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", required=True) + 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", + "process", + ] + ) + 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 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]) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/rerun/run_cleanup_benchmark.py b/benchmark/rerun/run_cleanup_benchmark.py new file mode 100644 index 00000000..a8e3ebe1 --- /dev/null +++ b/benchmark/rerun/run_cleanup_benchmark.py @@ -0,0 +1,142 @@ +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_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 f"{match.group('shard_id')}__w{match.group('worker_id')}" not in keep_keys: + paths_to_delete.add(rel_path) + return paths_to_delete + + +def _benchmark( + *, + mode: str, + fn: Callable[[list[str], set[str]], set[str]], + root_entries: list[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_keys)) + 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) + 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_keys=keep_keys, + iterations=args.iterations, + ), + _benchmark( + mode="fixed-slice", + fn=_cleanup_default_root_entries, + root_entries=root_entries, + keep_keys=keep_keys, + 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/benchmark/rerun/run_cloud_benchmark.py b/benchmark/rerun/run_cloud_benchmark.py new file mode 100644 index 00000000..99a7d517 --- /dev/null +++ b/benchmark/rerun/run_cloud_benchmark.py @@ -0,0 +1,638 @@ +from __future__ import annotations + +import argparse +import json +import os +import platform as platform_module +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 + job_error: str | None + 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 + stage_duration_s: float | None + stage_results: list[StageResult] + output_file_count: int | None + output_size_bytes: int | None + output_inspection_error: str | None + submitter_python_version: str + submitter_platform: str + git_ref: str + submitter_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.", + ) + parser.add_argument( + "--continue-on-failure", + action="store_true", + help="Continue running later cases after a cloud job fails.", + ) + 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", + materialize_tables=False, + ).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 _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 + + +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 + 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 + 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) + + +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, + 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, + ) + 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( + 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) + stage_results = _stage_results(client, job) + + return CaseResult( + case=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), + 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")), + 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, + submitter_python_version=sys.version.replace("\n", " "), + submitter_platform=platform_module.platform(), + git_ref=git_ref, + submitter_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 _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: + 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, + ) + 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_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 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmark/rerun/run_local_benchmark.py b/benchmark/rerun/run_local_benchmark.py new file mode 100644 index 00000000..ed4062fd --- /dev/null +++ b/benchmark/rerun/run_local_benchmark.py @@ -0,0 +1,216 @@ +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_ns +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_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( + 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("--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() + + +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, + writes_per_iteration: int, +) -> 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_ns() + if force_fallback: + with _force_chunk_fallback(): + for _ in range(writes_per_iteration): + sink.write_shard_block("shard-a", block) + else: + for _ in range(writes_per_iteration): + sink.write_shard_block("shard-a", block) + sink.on_shard_complete("shard-a") + wall_time_ns = perf_counter_ns() - start + written = sorted(output.glob("**/*.rrd")) + 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_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( + path.read_bytes() == source.read_bytes() for path in written + ), + ) + + +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, + writes_per_iteration=args.writes_per_iteration, + ) + results.append(result) + print( + f"{mode_name} iteration {iteration}: " + f"{result.wall_time_s:.6f}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):.6f}s " + "chunk-fallback avg=" + f"{sum(fallback) / len(fallback):.6f}s" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) 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..894df0a4 --- /dev/null +++ b/docs/reading-data/rerun.md @@ -0,0 +1,122 @@ +--- +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. + +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 + +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..566b3c75 --- /dev/null +++ b/docs/writing-data/rerun.md @@ -0,0 +1,71 @@ +--- +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. 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 + +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 +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 +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. 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 + +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. 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/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/io/datafile.py b/src/refiner/io/datafile.py index 827d2ca1..ae1b4208 100644 --- a/src/refiner/io/datafile.py +++ b/src/refiner/io/datafile.py @@ -146,6 +146,21 @@ 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, + block_size=8 * 1024 * 1024, + ) + 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/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/__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/_rerun_io.py b/src/refiner/pipeline/_rerun_io.py new file mode 100644 index 00000000..5388c364 --- /dev/null +++ b/src/refiner/pipeline/_rerun_io.py @@ -0,0 +1,82 @@ +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 + source_recording_count: int | None = None + + +__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 4ee33797..4cc0ae51 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, @@ -41,6 +41,7 @@ JsonReader, McapReader, ParquetReader, + RerunReader, TfdsReader, TfrecordReader, ZarrReader, @@ -48,6 +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 ( + 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 @@ -617,6 +624,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, @@ -710,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, @@ -729,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 @@ -746,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, @@ -1217,6 +1255,73 @@ 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, + materialize_tables: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + 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, + 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. 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( + 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, + materialize_tables=materialize_tables, + include_recording=include_recording, + fill_latest_at=fill_latest_at, + action_prefix=action_prefix, + state_prefix=state_prefix, + camera_prefix=camera_prefix, + actions=actions, + states=states, + videos=videos, + fps=fps, + robot_type=robot_type, + ) + ) + + def read_parquet( inputs: DataFileSetLike, *, 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/__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/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/reducer/file.py b/src/refiner/pipeline/sinks/reducer/file.py index f32b079f..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,23 +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) - } + 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] - ) - paths = [listing_prefix] - prefix_parts = [part for part in listing_prefix.split("/") if part] + 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(self._listing_prefix, detail=False) + except (FileNotFoundError, NotADirectoryError): + root_entries = [] + 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: + rm(rel_path, recursive=True) + except FileNotFoundError: + continue + return + 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: @@ -178,3 +201,16 @@ def _run_cleanup(self) -> None: __all__ = ["FileCleanupReducerSink"] + + +def _cleanup_default_root_entries( + root_entries: list[str], + keep_keys: set[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 + if rel_path not in keep_keys: + paths_to_delete.add(rel_path) + return paths_to_delete diff --git a/src/refiner/pipeline/sinks/rerun.py b/src/refiner/pipeline/sinks/rerun.py new file mode 100644 index 00000000..ef49270e --- /dev/null +++ b/src/refiner/pipeline/sinks/rerun.py @@ -0,0 +1,492 @@ +from __future__ import annotations + +import os +import shutil +import tempfile +import warnings +from pathlib import Path +from string import Formatter +from typing import Any, Callable + +import pyarrow as pa + +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.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" + + +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: + template_fields = _validate_filename_template(filename_template) + self.output = DataFolder.resolve(output) + self._local_output_root = ( + 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, + 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[str] = 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) + 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()) + ) + 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") + return count + + def _write_recording(self, recording: RerunRecording, relpath: str) -> None: + target = self.output.file(relpath) + + def write_local(path: Path | str) -> 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( + recording, + path, + application_id=self.app_id, + write_footer=self.write_footer, + ) + + local_output_root = self._local_output_root + if local_output_root is not None: + local_path = f"{local_output_root}/{relpath}" + self._ensure_local_parent(os.path.dirname(local_path)) + 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) + if not self._uses_row_index: + self._written_relpaths.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 _ensure_local_parent(self, parent: str) -> None: + if parent not in self._created_local_parents: + Path(parent).mkdir(parents=True, exist_ok=True) + self._created_local_parents.add(parent) + + +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 _write_source_chunks( + recording: RerunRecording, + path: Path | str, + *, + 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: + _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 LocalRrd(source) as local_path: + _write_source_chunks_from_path( + recording, + path, + local_path=local_path, + application_id=application_id, + ) + + +def _write_source_chunks_from_path( + recording: RerunRecording, + path: Path | str, + *, + local_path: Path, + application_id: str, +) -> None: + if _can_copy_source_rrd(recording): + try: + os.link(local_path, path) + except OSError: + shutil.copyfile(local_path, path) + return + + check_required_dependencies( + "write_rerun", + [("rerun", "rerun-sdk")], + dist="rerun", + ) + import rerun as rr + + 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: + 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 _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, + *, + 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 | str, + *, + application_id: str, + write_footer: bool, +) -> 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 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: + 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" + + +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 + ): + 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)) + ) + 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", + worker_id="worker", + row_index=0, + segment_id="segment", + ), + "filename_template", + ) + 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, + *, + shard_id: str, + worker_id: str, + row_index: int, + segment_id: str, + uses_segment_id: bool, +) -> str: + 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, + ) + + +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") + 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/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/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/__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/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 new file mode 100644 index 00000000..673de053 --- /dev/null +++ b/src/refiner/pipeline/sources/readers/rerun.py @@ -0,0 +1,969 @@ +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 +import warnings + +import numpy as np +import pyarrow as pa +import pyarrow.compute as pc +from fsspec import AbstractFileSystem + +from refiner.io import DataFile +from refiner.io.fileset import DataFileSetLike +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 +from refiner.pipeline.sources.readers.base import BaseReader +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 +from refiner.worker.context import logger + +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"} +) +_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 + + +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): + """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, + materialize_tables: bool = True, + include_recording: bool | None = None, + fill_latest_at: bool = False, + 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, + 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 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 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: + 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.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 + ) + 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.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") + 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": + 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) + 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 + + 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, + "include_static": self.include_static, + "materialize_tables": self.materialize_tables, + "include_recording": self.include_recording, + "fill_latest_at": self.fill_latest_at, + } + ) + 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) + batch: list[tuple[DataFile, 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) + 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, LocalRrd]], + ) -> Iterator[SourceUnit]: + opened_files = _open_local_sources(local_files) + if self._retain_batch_local_sources(): + try: + units = list(self._read_files(opened_files)) + except BaseException: + _close_local_sources(opened_files) + raise + if not units: + _close_local_sources(opened_files) + return + try: + yield cast(list[Row], units) + finally: + _close_local_sources(opened_files) + return + + try: + yield from self._read_files(opened_files) + finally: + _close_local_sources(opened_files) + + 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 ( + 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) + + def _read_files_with_server( + self, + local_files: Sequence[tuple[DataFile, Path, LocalRrd]], + ) -> Iterator[SourceUnit]: + check_required_dependencies( + "read_rerun", + [("rerun", "rerun-sdk"), "datafusion"], + dist="rerun", + ) + import rerun as rr + + datasets = { + f"recording_{index}": (str(local_path),) + 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, local_source) in zip( + datasets, local_files, strict=True + ): + dataset = client.get_dataset(dataset_name) + 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, + store_entries: Sequence[Any] | None = None, + ) -> Iterator[SourceUnit]: + 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 + 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 + recording_id = store.recording_id if store is not None else segment_id + view = dataset.filter_segments([segment_id]) + if self.output == "robotics": + yield self._robotics_row( + view, + 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, + ) + else: + 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 + } + yield self._recording_row( + segment_id=segment_id, + source=source, + local_source=local_source, + tables=tables, + 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( + self, + source: DataFile, + local_source: LocalRrd, + store_entries: Sequence[Any], + ) -> list[DictRow]: + rows = [] + source_recording_count = len(store_entries) + for store in store_entries: + recording_id = str(store.recording_id) + rows.append( + self._recording_row( + segment_id=recording_id, + source=source, + local_source=local_source, + tables={}, + static=None, + application_id=store.application_id, + recording_id=recording_id, + source_recording_count=source_recording_count, + ) + ) + return rows + + def _recording_row( + self, + *, + segment_id: str, + source: DataFile, + local_source: LocalRrd, + tables: Mapping[str, Tabular], + 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, + "rerun": RerunRecording( + segment_id=segment_id, + source_path=source.abs_path(), + 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, + source_recording_count=source_recording_count, + 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, ...]: + 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_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, + *, + segment_id: str, + source_path: str, + source_file: DataFile, + local_source: LocalRrd, + 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( + content_view.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: + static = ( + _collect_table(content_view.reader(index=None)) + if self.include_static + else None + ) + 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, + 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 + if self.robot_type is not None: + row["robot_type"] = self.robot_type + + component_columns = _component_column_maps(table) + scalar_columns = component_columns.get("Scalars:scalars", {}) + 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", + _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) + + image_columns = component_columns.get("EncodedImage:blob", {}) + camera_columns = ( + _selected_camera_columns(image_columns, self.videos) + 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() + _require_dense_encoded_images(values, video_name=name) + 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 _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: + 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") + 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: + table = df.to_arrow_table() + if _RERUN_SEGMENT_ID in table.column_names and table.num_rows > 0: + segment_ids = table.column(_RERUN_SEGMENT_ID) + if segment_ids.null_count: + table = table.filter(_is_valid(segment_ids)) + 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 _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 _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 + + try: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", + message="RRD file has no footer/manifest:.*", + ) + 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: {}", + type(err).__name__, + ) + return [] + + +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_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 {} + 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: + 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]]: + return sorted( + ( + (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(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 + 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 + + +def _selected_camera_columns( + by_entity_path: Mapping[str, str], + selected: Mapping[str, str], +) -> dict[str, str]: + 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 _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: + 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 = 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: + 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 _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 + 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 + values = np.asarray(array.values) + 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 + + from PIL import Image + + 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}" + ) + outer_offsets = np.asarray(values.offsets) + valid = np.asarray(_is_valid(values), dtype=bool) + inner = values.values + 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: + 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__ = [ + "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..82e0f493 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: @@ -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}" ) @@ -252,16 +268,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/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/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/io/test_datafile_datafolder.py b/tests/io/test_datafile_datafolder.py index e16849ad..73a8fdeb 100644 --- a/tests/io/test_datafile_datafolder.py +++ b/tests/io/test_datafile_datafolder.py @@ -85,6 +85,33 @@ 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, dict[str, Any]]] = [] + original_put_file = fs.put_file + + def fake_put_file(lpath, rpath, **kwargs): + 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", + {"block_size": 8 * 1024 * 1024}, + ) + ] + assert fs.cat("bucket/dest.txt") == b"payload" + + def test_datafile_resolve_adds_hf_token_for_huggingface_http_urls(monkeypatch): captured = {} 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/pipeline/test_sinks.py b/tests/pipeline/test_sinks.py index 07ef0d55..d2a11b80 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 @@ -964,6 +965,66 @@ 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_cleanup_default_root_entries_returns_only_losers() -> None: + root_entries = [ + "0123456789ab__w111111111111", + "0123456789ab__w222222222222", + "not-a-match", + ] + keep_keys = {"0123456789ab__w111111111111"} + + assert _cleanup_default_root_entries(root_entries, keep_keys) == { + "0123456789ab__w222222222222" + } + + def test_file_cleanup_reducer_removes_non_finalized_nested_directories( tmp_path, ) -> 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..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 @@ -45,6 +46,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 +269,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: @@ -426,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()), @@ -435,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, @@ -452,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 new file mode 100644 index 00000000..fa16503a --- /dev/null +++ b/tests/readers/test_rerun_reader.py @@ -0,0 +1,1028 @@ +from __future__ import annotations + +from dataclasses import replace +from io import BytesIO +from pathlib import Path +from typing import Any, cast + +import fsspec +import numpy as np +import pytest + +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 RerunReader, RerunRecording + +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( + "/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)], + columns=rr.Scalars.columns(scalars=np.asarray([4.0, 5.0, 6.0])), + ) + + +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 _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 _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 + + 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 _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) + + 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 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: + 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) + + 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_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",) + 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) + + 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) + 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( + 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_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) + + 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]] + + +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_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_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_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_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: + 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_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, +) -> 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" + _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" + _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 + + +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) + + 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 + 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( + "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") + 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() + assert recording.source_recording_count == 1 + + 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", 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 + 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( + 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_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_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_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: + 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: + 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_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, +) -> 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_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) + + 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 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 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"