#!/usr/bin/env python3 """Rebuild filtered T-Rex LeRobot-v2 base/Track-Force dataset variants. The builder is resumable and writes only under a staging directory until the explicit install phase. RGB/tactile videos are hard-linked; Parquet files and renumbered Track NPZ files are rewritten atomically. """ from __future__ import annotations import argparse import copy import hashlib import json import math import os import shutil import sys import tempfile from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor, as_completed from dataclasses import asdict, dataclass from datetime import datetime, timezone from pathlib import Path from typing import Iterable, Sequence import numpy as np _DREAMZERO_ROOT = Path(__file__).resolve().parents[2] _SCRIPTS_ROOT = _DREAMZERO_ROOT / "scripts" for _path in (_SCRIPTS_ROOT, _SCRIPTS_ROOT / "data"): if str(_path) not in sys.path: sys.path.insert(0, str(_path)) import build_trex_track_force_v2 as force_builder # noqa: E402 TRACK_CACHE_NAME = "tracks_trex_track_force_v2" VARIANT_NAMES = ( "trex_small", "trex_small_force", "trex_full", "trex_full_force", ) LEGACY_NAMES = ("trex_small", "trex_datasetv2", "trex_track_force_v2") PARQUET_SCHEMA_METADATA_KEY = b"trex_track_force_schema_version" @dataclass(frozen=True) class EpisodeRecord: source_episode_index: int episode_index: int length: int task: str task_index: int global_index_start: int def _utc_now() -> str: return datetime.now(timezone.utc).isoformat() def _read_json(path: Path) -> dict: return json.loads(path.read_text()) def _read_jsonl(path: Path) -> list[dict]: return [ json.loads(line) for line in path.read_text().splitlines() if line.strip() ] def _atomic_write_json(path: Path, value: object) -> None: path.parent.mkdir(parents=True, exist_ok=True) fd, temporary_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, ) try: with os.fdopen(fd, "w") as file: json.dump(value, file, indent=2) file.write("\n") os.replace(temporary_name, path) except Exception: Path(temporary_name).unlink(missing_ok=True) raise def _atomic_write_jsonl(path: Path, rows: Iterable[dict]) -> None: path.parent.mkdir(parents=True, exist_ok=True) fd, temporary_name = tempfile.mkstemp( prefix=f".{path.name}.", suffix=".tmp", dir=path.parent, ) try: with os.fdopen(fd, "w") as file: for row in rows: file.write(json.dumps(row) + "\n") os.replace(temporary_name, path) except Exception: Path(temporary_name).unlink(missing_ok=True) raise def _sha256(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as file: while chunk := file.read(8 * 1024 * 1024): digest.update(chunk) return digest.hexdigest() def _video_keys(info: dict) -> list[str]: return sorted( key for key, feature in info["features"].items() if feature.get("dtype") == "video" ) def _source_episode_path(root: Path, episode_index: int, suffix: str) -> Path: return ( root / suffix / f"chunk-{episode_index // 1000:03d}" / f"episode_{episode_index:06d}.parquet" ) def _episode_parquet_path(root: Path, episode_index: int) -> Path: return ( root / "data" / f"chunk-{episode_index // 1000:03d}" / f"episode_{episode_index:06d}.parquet" ) def _episode_video_path(root: Path, episode_index: int, video_key: str) -> Path: return ( root / "videos" / f"chunk-{episode_index // 1000:03d}" / video_key / f"episode_{episode_index:06d}.mp4" ) def _track_path(root: Path, episode_index: int) -> Path: return root / TRACK_CACHE_NAME / f"episode_{episode_index:06d}.npz" def _build_registry( episodes: Sequence[dict], source_tasks: Sequence[dict], *, excluded: set[int], source_limit: int | None, ) -> tuple[list[EpisodeRecord], list[dict]]: tasks_by_text = { str(row["task"]): int(row["task_index"]) for row in source_tasks } selected = [ episode for episode in episodes if int(episode["episode_index"]) not in excluded and ( source_limit is None or int(episode["episode_index"]) < int(source_limit) ) ] retained_tasks = { str(task) for episode in selected for task in episode["tasks"] } missing_tasks = retained_tasks.difference(tasks_by_text) if missing_tasks: raise ValueError(f"episodes reference unknown tasks: {sorted(missing_tasks)[:3]}") ordered_tasks = sorted(retained_tasks, key=tasks_by_text.__getitem__) new_task_by_text = { task: task_index for task_index, task in enumerate(ordered_tasks) } task_rows = [ {"task_index": task_index, "task": task} for task, task_index in new_task_by_text.items() ] registry: list[EpisodeRecord] = [] global_index_start = 0 for new_episode_index, episode in enumerate(selected): tasks = [str(task) for task in episode["tasks"]] if len(tasks) != 1: raise ValueError( f"episode {episode['episode_index']} has {len(tasks)} tasks; expected one" ) length = int(episode["length"]) registry.append( EpisodeRecord( source_episode_index=int(episode["episode_index"]), episode_index=new_episode_index, length=length, task=tasks[0], task_index=new_task_by_text[tasks[0]], global_index_start=global_index_start, ) ) global_index_start += length return registry, task_rows def _replace_primitive_column(table, name: str, values: np.ndarray): import pyarrow as pa index = table.schema.get_field_index(name) if index < 0: raise ValueError(f"missing required parquet column {name!r}") field = table.schema.field(index) return table.set_column(index, field, pa.array(values, type=field.type)) def _parquet_is_ready( path: Path, *, episode_index: int, task_index: int, global_index_start: int, expected_rows: int, force: bool, ) -> bool: if not path.is_file(): return False try: import pyarrow.parquet as pq parquet_file = pq.ParquetFile(path) if int(parquet_file.metadata.num_rows) != int(expected_rows): return False metadata = parquet_file.schema_arrow.metadata or {} if force and metadata.get(PARQUET_SCHEMA_METADATA_KEY) != ( force_builder.SCHEMA_VERSION.encode() ): return False table = pq.read_table( path, columns=["episode_index", "task_index", "frame_index", "index"], ) stored_episode = np.asarray(table["episode_index"].to_numpy()) stored_task = np.asarray(table["task_index"].to_numpy()) frame_index = np.asarray(table["frame_index"].to_numpy()) global_index = np.asarray(table["index"].to_numpy()) return bool( np.all(stored_episode == int(episode_index)) and np.all(stored_task == int(task_index)) and np.array_equal(frame_index, np.arange(expected_rows)) and np.array_equal( global_index, np.arange( global_index_start, global_index_start + expected_rows, ), ) ) except Exception: return False def _rewrite_parquet_job(job: tuple) -> tuple[int, str]: ( source_path_text, destination_path_text, episode_index, task_index, global_index_start, expected_rows, force, ) = job source_path = Path(source_path_text) destination_path = Path(destination_path_text) if _parquet_is_ready( destination_path, episode_index=episode_index, task_index=task_index, global_index_start=global_index_start, expected_rows=expected_rows, force=force, ): return episode_index, "skipped" import pyarrow.parquet as pq table = pq.read_table(source_path) if int(table.num_rows) != int(expected_rows): raise ValueError( f"{source_path}: {table.num_rows} rows != expected {expected_rows}" ) original_metadata = table.schema.metadata table = _replace_primitive_column( table, "episode_index", np.full(expected_rows, episode_index, dtype=np.int64), ) table = _replace_primitive_column( table, "task_index", np.full(expected_rows, task_index, dtype=np.int64), ) table = _replace_primitive_column( table, "frame_index", np.arange(expected_rows, dtype=np.int64), ) table = _replace_primitive_column( table, "index", np.arange( global_index_start, global_index_start + expected_rows, dtype=np.int64, ), ) table = table.replace_schema_metadata(original_metadata) if force: metadata = table.schema.metadata or {} if metadata.get(PARQUET_SCHEMA_METADATA_KEY) != ( force_builder.SCHEMA_VERSION.encode() ): raise ValueError(f"{source_path}: force schema metadata is missing") destination_path.parent.mkdir(parents=True, exist_ok=True) temporary_path = destination_path.with_suffix(".tmp.parquet") temporary_path.unlink(missing_ok=True) try: pq.write_table( table, temporary_path, compression="zstd", use_dictionary=True, ) os.replace(temporary_path, destination_path) except Exception: temporary_path.unlink(missing_ok=True) raise return episode_index, "written" def _track_is_ready( path: Path, *, episode_index: int, expected_frames: int, ) -> bool: if not path.is_file(): return False try: with np.load(path, allow_pickle=False) as payload: return bool( int(np.asarray(payload["episode_index"]).item()) == int(episode_index) and int(np.asarray(payload["num_steps"]).item()) == int(expected_frames) ) except Exception: return False def _rewrite_track_job(job: tuple) -> tuple[int, str, str]: ( source_path_text, destination_path_text, source_episode_index, episode_index, expected_frames, ) = job source_path = Path(source_path_text) destination_path = Path(destination_path_text) if _track_is_ready( destination_path, episode_index=episode_index, expected_frames=expected_frames, ): return episode_index, _sha256(destination_path), "skipped" destination_path.parent.mkdir(parents=True, exist_ok=True) if int(source_episode_index) == int(episode_index): destination_path.unlink(missing_ok=True) os.link(source_path, destination_path) return episode_index, _sha256(destination_path), "linked" with np.load(source_path, allow_pickle=False) as archive: payload = { name: np.asarray(archive[name]).copy() for name in archive.files } payload["episode_index"] = np.array(episode_index, dtype=np.int32) payload["source_episode_index"] = np.array( source_episode_index, dtype=np.int32, ) temporary_path = destination_path.with_suffix(".tmp.npz") temporary_path.unlink(missing_ok=True) try: with temporary_path.open("wb") as file: np.savez_compressed(file, **payload) os.replace(temporary_path, destination_path) except Exception: temporary_path.unlink(missing_ok=True) raise return episode_index, _sha256(destination_path), "written" def _hardlink(source: Path, destination: Path) -> str: if destination.exists(): if os.path.samefile(source, destination): return "skipped" raise FileExistsError(f"destination is not the expected hard link: {destination}") destination.parent.mkdir(parents=True, exist_ok=True) os.link(source, destination) return "linked" def _link_job(job: tuple[str, str]) -> str: return _hardlink(Path(job[0]), Path(job[1])) def _run_jobs( jobs: Sequence[tuple], worker, *, workers: int, process: bool, label: str, ) -> list: if not jobs: return [] executor_type = ProcessPoolExecutor if process else ThreadPoolExecutor results = [] with executor_type(max_workers=max(1, workers)) as executor: futures = [executor.submit(worker, job) for job in jobs] for completed, future in enumerate(as_completed(futures), start=1): results.append(future.result()) if completed % 250 == 0 or completed == len(futures): print(f"{label}: {completed}/{len(futures)}", flush=True) return results def _write_variant_metadata_skeleton( *, variant_root: Path, final_root: Path, source_info: dict, source_modality: dict, source_embodiment: dict, registry: Sequence[EpisodeRecord], tasks: Sequence[dict], video_keys: Sequence[str], excluded_source_indices: Sequence[int], force: bool, source_dataset: Path, blacklist_sha256: str, ) -> None: info = copy.deepcopy(source_info) info["total_episodes"] = len(registry) info["total_frames"] = int(sum(record.length for record in registry)) info["total_tasks"] = len(tasks) info["total_videos"] = len(registry) * len(video_keys) info["total_chunks"] = math.ceil(len(registry) / int(info["chunks_size"])) info["splits"] = {"train": f"0:{len(registry)}"} info.pop("discarded_episode_indices", None) info.pop("trex_track_force", None) info["trex_filter"] = { "source_dataset": str(source_dataset), "source_episode_count": int(source_info["total_episodes"]), "source_to_new_map": "meta/source_episode_index_map.json", "excluded_source_indices": "meta/excluded_source_episode_indices.json", "blacklist_sha256": blacklist_sha256, "force_variant": bool(force), "created_at": _utc_now(), } metadata_dir = variant_root / "meta" metadata_dir.mkdir(parents=True, exist_ok=True) _atomic_write_json(metadata_dir / "info.json", info) _atomic_write_json(metadata_dir / "modality.json", source_modality) _atomic_write_json(metadata_dir / "embodiment.json", source_embodiment) _atomic_write_jsonl( metadata_dir / "episodes.jsonl", ( { "episode_index": record.episode_index, "tasks": [record.task], "length": record.length, } for record in registry ), ) _atomic_write_jsonl(metadata_dir / "tasks.jsonl", tasks) _atomic_write_json( metadata_dir / "source_episode_index_map.json", { str(record.source_episode_index): record.episode_index for record in registry }, ) _atomic_write_json( metadata_dir / "episode_index_provenance.json", [asdict(record) for record in registry], ) _atomic_write_json( metadata_dir / "excluded_source_episode_indices.json", list(excluded_source_indices), ) _atomic_write_json( metadata_dir / "dataset_variant.json", { "name": final_root.name, "final_root": str(final_root), "force": bool(force), "episodes": len(registry), "frames": int(sum(record.length for record in registry)), "tasks": len(tasks), "video_keys": list(video_keys), "created_at": _utc_now(), }, ) def _rewrite_parquets( *, source_root: Path, destination_root: Path, registry: Sequence[EpisodeRecord], workers: int, force: bool, ) -> None: jobs = [ ( str(_episode_parquet_path(source_root, record.source_episode_index)), str(_episode_parquet_path(destination_root, record.episode_index)), record.episode_index, record.task_index, record.global_index_start, record.length, force, ) for record in registry ] _run_jobs( jobs, _rewrite_parquet_job, workers=workers, process=True, label=f"{destination_root.name} parquet", ) def _link_videos( *, source_root: Path, destination_root: Path, registry: Sequence[EpisodeRecord], video_keys: Sequence[str], source_uses_new_indices: bool, workers: int, ) -> None: jobs: list[tuple[str, str]] = [] for record in registry: source_episode_index = ( record.episode_index if source_uses_new_indices else record.source_episode_index ) for video_key in video_keys: jobs.append( ( str( _episode_video_path( source_root, source_episode_index, video_key, ) ), str( _episode_video_path( destination_root, record.episode_index, video_key, ) ), ) ) _run_jobs( jobs, _link_job, workers=workers, process=False, label=f"{destination_root.name} videos", ) def _rewrite_tracks( *, source_root: Path, destination_root: Path, registry: Sequence[EpisodeRecord], workers: int, ) -> dict[int, str]: jobs = [ ( str(_track_path(source_root, record.source_episode_index)), str(_track_path(destination_root, record.episode_index)), record.source_episode_index, record.episode_index, record.length, ) for record in registry ] results = _run_jobs( jobs, _rewrite_track_job, workers=workers, process=True, label=f"{destination_root.name} tracks", ) return {int(episode_index): sha256 for episode_index, sha256, _ in results} def _link_subset_files( *, source_root: Path, destination_root: Path, registry: Sequence[EpisodeRecord], video_keys: Sequence[str], force: bool, workers: int, ) -> dict[int, str]: parquet_jobs = [ ( str(_episode_parquet_path(source_root, record.episode_index)), str(_episode_parquet_path(destination_root, record.episode_index)), ) for record in registry ] _run_jobs( parquet_jobs, _link_job, workers=workers, process=False, label=f"{destination_root.name} parquet links", ) _link_videos( source_root=source_root, destination_root=destination_root, registry=registry, video_keys=video_keys, source_uses_new_indices=True, workers=workers, ) track_hashes: dict[int, str] = {} if force: track_jobs = [ ( str(_track_path(source_root, record.episode_index)), str(_track_path(destination_root, record.episode_index)), ) for record in registry ] _run_jobs( track_jobs, _link_job, workers=workers, process=False, label=f"{destination_root.name} track links", ) track_hashes = { record.episode_index: _sha256( _track_path(destination_root, record.episode_index) ) for record in registry } return track_hashes def _aggregate_source_manifest_entries(source_force_root: Path) -> dict[int, dict]: manifests = sorted( ( source_force_root / "meta" / "trex_track_force_prepare" ).glob("run_*/task_*.json"), key=lambda path: path.stat().st_mtime_ns, ) entries: dict[int, dict] = {} for path in manifests: try: payload = _read_json(path) except Exception: continue for key, entry in payload.get("episodes", {}).items(): if entry.get("status") == "complete": entries[int(key)] = copy.deepcopy(entry) expected = int( _read_json(source_force_root / "meta" / "info.json")["total_episodes"] ) missing = [index for index in range(expected) if index not in entries] if missing: raise ValueError( f"source task manifests do not cover every episode: {missing[:10]}" ) return entries def _build_force_manifest( *, variant_root: Path, final_root: Path, registry: Sequence[EpisodeRecord], source_entries: dict[int, dict], track_hashes: dict[int, str], ) -> dict: final_track_cache = final_root / TRACK_CACHE_NAME manifest = force_builder._new_manifest(final_root, final_track_cache) manifest["filter_provenance"] = { "source_to_new_map": "meta/source_episode_index_map.json", "excluded_source_indices": "meta/excluded_source_episode_indices.json", } for record in registry: entry = copy.deepcopy(source_entries[record.source_episode_index]) entry["episode_index"] = record.episode_index entry["source_episode_index"] = record.source_episode_index entry["num_frames"] = record.length entry["parquet"] = str( _episode_parquet_path(Path("."), record.episode_index) ) entry["track_npz"] = str( final_track_cache / f"episode_{record.episode_index:06d}.npz" ) entry["track_sha256"] = track_hashes[record.episode_index] entry["status"] = "complete" entry["skipped"] = False manifest["episodes"][f"{record.episode_index:06d}"] = entry _atomic_write_json( variant_root / "meta" / "trex_track_force_manifest.json", manifest, ) return manifest def _compute_base_stats(dataset_root: Path, registry: Sequence[EpisodeRecord]) -> dict: import pyarrow.parquet as pq stats: dict[str, dict] = {} for column in ("observation.state", "action", "timestamp"): parts: list[np.ndarray] = [] for record in registry: table = pq.read_table( _episode_parquet_path(dataset_root, record.episode_index), columns=[column], ) values = force_builder._column_to_numpy(table, column) if values.ndim == 1: values = values[:, None] parts.append(values) stats[column] = force_builder._statistics(np.concatenate(parts, axis=0)) return stats def _compute_base_relative_stats( dataset_root: Path, registry: Sequence[EpisodeRecord], ) -> dict: import pyarrow.parquet as pq action_offsets = range(24) output: dict[str, dict] = {} for name, selection in ( ("left_arm", slice(0, 7)), ("right_arm", slice(29, 36)), ): parts: list[np.ndarray] = [] for record in registry: table = pq.read_table( _episode_parquet_path(dataset_root, record.episode_index), columns=["observation.state", "action"], ) state = force_builder._column_to_numpy( table, "observation.state", )[:, selection] action = force_builder._column_to_numpy( table, "action", )[:, selection] usable_length = len(state) - max(action_offsets) if usable_length <= 0: continue reference = state[:usable_length] parts.extend( action[offset : offset + usable_length] - reference for offset in action_offsets ) if not parts: raise ValueError(f"no relative action samples for {name}") output[name] = force_builder._statistics( np.concatenate(parts, axis=0) ) return output def _finalize_metadata( *, variant_root: Path, registry: Sequence[EpisodeRecord], force: bool, ) -> None: stats_path = variant_root / "meta" / "stats.json" if force: base_variant_root = variant_root.parent / variant_root.name.removesuffix( "_force" ) base_stats_path = base_variant_root / "meta" / "stats.json" if not base_stats_path.is_file(): raise FileNotFoundError( f"base variant stats must be finalized first: {base_stats_path}" ) _atomic_write_json(stats_path, _read_json(base_stats_path)) else: _atomic_write_json( stats_path, _compute_base_stats(variant_root, registry), ) if not force: _atomic_write_json( variant_root / "meta" / "relative_stats_dreamzero.json", _compute_base_relative_stats(variant_root, registry), ) return result = force_builder.update_metadata( variant_root, assume_all_converted=True, ) manifest_path = ( variant_root / "meta" / "trex_track_force_manifest.json" ) manifest = _read_json(manifest_path) manifest["metadata"] = result manifest["updated_at"] = _utc_now() _atomic_write_json(manifest_path, manifest) force_builder.validate_metadata(variant_root) for backup in (variant_root / "meta").glob("*.trex_track_force.bak"): backup.unlink() def _copy_audit(source_force_root: Path, destination_force_root: Path) -> None: source = source_force_root / "audit" destination = destination_force_root / "audit" / "source_dataset" if destination.exists() or not source.exists(): return shutil.copytree(source, destination) def _validate_variant( *, root: Path, expected_registry: Sequence[EpisodeRecord], video_keys: Sequence[str], force: bool, ) -> dict: import pyarrow.parquet as pq info = _read_json(root / "meta" / "info.json") episodes = _read_jsonl(root / "meta" / "episodes.jsonl") tasks = _read_jsonl(root / "meta" / "tasks.jsonl") if int(info["total_episodes"]) != len(expected_registry): raise ValueError(f"{root}: total_episodes is stale") if int(info["total_frames"]) != sum(r.length for r in expected_registry): raise ValueError(f"{root}: total_frames is stale") if int(info["total_tasks"]) != len(tasks): raise ValueError(f"{root}: total_tasks is stale") if int(info["total_videos"]) != len(expected_registry) * len(video_keys): raise ValueError(f"{root}: total_videos is stale") if [int(row["episode_index"]) for row in episodes] != list( range(len(expected_registry)) ): raise ValueError(f"{root}: episodes.jsonl is not dense") if [int(row["task_index"]) for row in tasks] != list(range(len(tasks))): raise ValueError(f"{root}: tasks.jsonl is not dense") manifest = None if force: force_builder.validate_metadata(root) manifest = _read_json( root / "meta" / "trex_track_force_manifest.json" ) if len(manifest.get("episodes", {})) != len(expected_registry): raise ValueError(f"{root}: force manifest count is stale") for completed, record in enumerate(expected_registry, start=1): path = _episode_parquet_path(root, record.episode_index) if not _parquet_is_ready( path, episode_index=record.episode_index, task_index=record.task_index, global_index_start=record.global_index_start, expected_rows=record.length, force=force, ): raise ValueError(f"{root}: invalid parquet {path}") for video_key in video_keys: video_path = _episode_video_path( root, record.episode_index, video_key, ) if not video_path.is_file() or video_path.stat().st_size <= 0: raise ValueError(f"{root}: missing video {video_path}") if force: track_path = _track_path(root, record.episode_index) if not _track_is_ready( track_path, episode_index=record.episode_index, expected_frames=record.length, ): raise ValueError(f"{root}: invalid track cache {track_path}") entry = manifest["episodes"].get( f"{record.episode_index:06d}", {}, ) if ( entry.get("status") != "complete" or int(entry.get("source_episode_index", -1)) != record.source_episode_index ): raise ValueError( f"{root}: invalid manifest entry {record.episode_index}" ) if completed % 250 == 0 or completed == len(expected_registry): print( f"validate {root.name}: {completed}/{len(expected_registry)}", flush=True, ) sample_indices = sorted( { 0, len(expected_registry) // 2, len(expected_registry) - 1, } ) for episode_index in sample_indices: parquet_file = pq.ParquetFile( _episode_parquet_path(root, episode_index) ) if int(parquet_file.metadata.num_rows) <= 0: raise ValueError(f"{root}: empty sample parquet {episode_index}") if force: force_builder.validate_episode_parquet( _episode_parquet_path(root, episode_index), verify_source_fk=False, ) result = { "root": str(root), "episodes": len(expected_registry), "frames": int(sum(record.length for record in expected_registry)), "tasks": len(tasks), "videos": len(expected_registry) * len(video_keys), "force": force, "validated_at": _utc_now(), } _atomic_write_json(root / "meta" / "dataset_ready.json", result) return result def _stage_paths(stage_root: Path) -> dict[str, Path]: return {name: stage_root / name for name in VARIANT_NAMES} def _load_inputs(args: argparse.Namespace): source_info = _read_json(args.base_root / "meta" / "info.json") source_modality = _read_json(args.base_root / "meta" / "modality.json") source_embodiment = _read_json(args.base_root / "meta" / "embodiment.json") episodes = _read_jsonl(args.base_root / "meta" / "episodes.jsonl") source_tasks = _read_jsonl(args.base_root / "meta" / "tasks.jsonl") excluded = set(json.loads(args.blacklist.read_text())) if len(excluded) != 115: raise ValueError( f"expected 115 excluded episodes, got {len(excluded)}" ) full_registry, full_tasks = _build_registry( episodes, source_tasks, excluded=excluded, source_limit=None, ) small_registry, small_tasks = _build_registry( episodes, source_tasks, excluded=excluded, source_limit=100, ) if len(full_registry) != 5349 or len(small_registry) != 100: raise ValueError( f"unexpected registry sizes: full={len(full_registry)}, " f"small={len(small_registry)}" ) return ( source_info, source_modality, source_embodiment, episodes, source_tasks, excluded, full_registry, full_tasks, small_registry, small_tasks, ) def _build(args: argparse.Namespace) -> None: ( source_info, source_modality, source_embodiment, _, _, excluded, full_registry, full_tasks, small_registry, small_tasks, ) = _load_inputs(args) paths = _stage_paths(args.stage_root) video_keys = _video_keys(source_info) blacklist_sha256 = _sha256(args.blacklist) final_paths = { name: args.data_root / name for name in VARIANT_NAMES } args.stage_root.mkdir(parents=True, exist_ok=True) for name, registry, tasks, force, source_dataset in ( ( "trex_full", full_registry, full_tasks, False, args.base_root, ), ( "trex_full_force", full_registry, full_tasks, True, args.force_root, ), ( "trex_small", small_registry, small_tasks, False, args.base_root, ), ( "trex_small_force", small_registry, small_tasks, True, args.force_root, ), ): _write_variant_metadata_skeleton( variant_root=paths[name], final_root=final_paths[name], source_info=source_info, source_modality=source_modality, source_embodiment=source_embodiment, registry=registry, tasks=tasks, video_keys=video_keys, excluded_source_indices=( sorted(excluded) if "full" in name else [] ), force=force, source_dataset=source_dataset, blacklist_sha256=blacklist_sha256, ) _rewrite_parquets( source_root=args.base_root, destination_root=paths["trex_full"], registry=full_registry, workers=args.parquet_workers, force=False, ) _link_videos( source_root=args.base_root, destination_root=paths["trex_full"], registry=full_registry, video_keys=video_keys, source_uses_new_indices=False, workers=args.link_workers, ) _rewrite_parquets( source_root=args.force_root, destination_root=paths["trex_full_force"], registry=full_registry, workers=args.parquet_workers, force=True, ) _link_videos( source_root=paths["trex_full"], destination_root=paths["trex_full_force"], registry=full_registry, video_keys=video_keys, source_uses_new_indices=True, workers=args.link_workers, ) full_track_hashes = _rewrite_tracks( source_root=args.force_root, destination_root=paths["trex_full_force"], registry=full_registry, workers=args.track_workers, ) source_entries = _aggregate_source_manifest_entries(args.force_root) _build_force_manifest( variant_root=paths["trex_full_force"], final_root=final_paths["trex_full_force"], registry=full_registry, source_entries=source_entries, track_hashes=full_track_hashes, ) _copy_audit(args.force_root, paths["trex_full_force"]) for small_record, full_record in zip( small_registry, full_registry[: len(small_registry)], ): if ( small_record.source_episode_index != full_record.source_episode_index or small_record.episode_index != full_record.episode_index or small_record.task_index != full_record.task_index or small_record.global_index_start != full_record.global_index_start ): raise ValueError("small is not an identity prefix of full") _link_subset_files( source_root=paths["trex_full"], destination_root=paths["trex_small"], registry=small_registry, video_keys=video_keys, force=False, workers=args.link_workers, ) small_track_hashes = _link_subset_files( source_root=paths["trex_full_force"], destination_root=paths["trex_small_force"], registry=small_registry, video_keys=video_keys, force=True, workers=args.link_workers, ) _build_force_manifest( variant_root=paths["trex_small_force"], final_root=final_paths["trex_small_force"], registry=small_registry, source_entries=source_entries, track_hashes=small_track_hashes, ) _atomic_write_json( args.stage_root / "build_complete.json", { "completed_at": _utc_now(), "variants": list(VARIANT_NAMES), }, ) def _metadata(args: argparse.Namespace) -> None: ( _, _, _, _, _, _, full_registry, _, small_registry, _, ) = _load_inputs(args) paths = _stage_paths(args.stage_root) selected = set(args.metadata_variants) for name, registry, force in ( ("trex_small", small_registry, False), ("trex_small_force", small_registry, True), ("trex_full", full_registry, False), ("trex_full_force", full_registry, True), ): if name not in selected: continue print(f"Finalizing metadata: {name}", flush=True) _finalize_metadata( variant_root=paths[name], registry=registry, force=force, ) _atomic_write_json( args.stage_root / "metadata_complete.json", {"completed_at": _utc_now()}, ) def _validate(args: argparse.Namespace) -> None: ( source_info, _, _, _, _, _, full_registry, _, small_registry, _, ) = _load_inputs(args) paths = _stage_paths(args.stage_root) video_keys = _video_keys(source_info) results = [] for name, registry, force in ( ("trex_small", small_registry, False), ("trex_small_force", small_registry, True), ("trex_full", full_registry, False), ("trex_full_force", full_registry, True), ): results.append( _validate_variant( root=paths[name], expected_registry=registry, video_keys=video_keys, force=force, ) ) _atomic_write_json( args.stage_root / "validation_complete.json", { "validated_at": _utc_now(), "results": results, }, ) def _install(args: argparse.Namespace) -> None: validation_marker = args.stage_root / "validation_complete.json" if not validation_marker.is_file(): raise RuntimeError("staging validation marker is missing") paths = _stage_paths(args.stage_root) if any(not path.is_dir() for path in paths.values()): raise RuntimeError("one or more staged variants are missing") backup_root = args.data_root / ( f".trex_old_{datetime.now().strftime('%Y%m%dT%H%M%S')}" ) backup_root.mkdir(parents=True, exist_ok=False) moved_old: list[tuple[Path, Path]] = [] installed: list[tuple[Path, Path]] = [] try: for name in dict.fromkeys((*LEGACY_NAMES, *VARIANT_NAMES)): current = args.data_root / name if current.exists(): backup = backup_root / name os.replace(current, backup) moved_old.append((backup, current)) for name in VARIANT_NAMES: staged = paths[name] final = args.data_root / name os.replace(staged, final) installed.append((final, staged)) ( source_info, _, _, _, _, _, full_registry, _, small_registry, _, ) = _load_inputs_from_backups(args, backup_root) video_keys = _video_keys(source_info) for name, registry, force in ( ("trex_small", small_registry, False), ("trex_small_force", small_registry, True), ("trex_full", full_registry, False), ("trex_full_force", full_registry, True), ): _validate_variant( root=args.data_root / name, expected_registry=registry, video_keys=video_keys, force=force, ) except Exception: for final, staged in reversed(installed): if final.exists(): os.replace(final, staged) for backup, current in reversed(moved_old): if backup.exists(): os.replace(backup, current) backup_root.rmdir() raise shutil.rmtree(backup_root) args.stage_root.mkdir(parents=True, exist_ok=True) for marker in args.stage_root.glob("*.json"): marker.unlink() args.stage_root.rmdir() def _load_inputs_from_backups(args: argparse.Namespace, backup_root: Path): backup_args = copy.copy(args) backup_args.base_root = backup_root / "trex_datasetv2" backup_args.force_root = backup_root / "trex_track_force_v2" backup_args.blacklist = ( backup_args.force_root / "audit" / "track_quality" / "frozen_wrist_episode_indices.json" ) return _load_inputs(backup_args) def _build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--data-root", type=Path, default=_DREAMZERO_ROOT / "data", ) parser.add_argument( "--base-root", type=Path, default=_DREAMZERO_ROOT / "data" / "trex_datasetv2", ) parser.add_argument( "--force-root", type=Path, default=_DREAMZERO_ROOT / "data" / "trex_track_force_v2", ) parser.add_argument( "--blacklist", type=Path, default=( _DREAMZERO_ROOT / "data" / "trex_track_force_v2" / "audit" / "track_quality" / "frozen_wrist_episode_indices.json" ), ) parser.add_argument( "--stage-root", type=Path, default=_DREAMZERO_ROOT / "data" / ".trex_variants_staging", ) parser.add_argument( "--phase", choices=("build", "metadata", "validate", "install", "all"), default="all", ) parser.add_argument("--parquet-workers", type=int, default=8) parser.add_argument("--track-workers", type=int, default=8) parser.add_argument("--link-workers", type=int, default=32) parser.add_argument( "--metadata-variants", nargs="+", choices=VARIANT_NAMES, default=list(VARIANT_NAMES), help="variants to process during the metadata phase", ) return parser def main(argv: Sequence[str] | None = None) -> int: args = _build_parser().parse_args(argv) for name in ("data_root", "base_root", "force_root", "blacklist", "stage_root"): setattr(args, name, getattr(args, name).expanduser().resolve()) if args.phase in ("build", "all"): _build(args) if args.phase in ("metadata", "all"): _metadata(args) if args.phase in ("validate", "all"): _validate(args) if args.phase in ("install", "all"): _install(args) return 0 if __name__ == "__main__": raise SystemExit(main())