"""Pinned LeRobot pipeline worker. Process selected shards and checkpoint to a new output repo.""" import hashlib import json import math import os import shutil import threading import time from collections import defaultdict from concurrent.futures import ThreadPoolExecutor, as_completed from contextvars import ContextVar from pathlib import Path import av import pyarrow.parquet as pq from huggingface_hub import CommitOperationAdd, HfApi, hf_hub_download, snapshot_download from human_video import FalMedia, generate_human_videos from prompt_config import validate_prompts from provider import Budget, BudgetExceededError, DeadlineExceededError, Provider from studio import COPY, FEATURES, HUMAN_VIDEOS, read_selection, relative_path from subset import EpisodeSubset from lerobot.annotations.steerable_pipeline.config import InterjectionsConfig, PlanConfig, VqaConfig from lerobot.annotations.steerable_pipeline.executor import Executor from lerobot.annotations.steerable_pipeline.modules import ( GeneralVqaModule, InterjectionsAndSpeechModule, general_vqa as vqa_module, interjections_and_speech as interjections_module, plan_subtasks_memory as plan_module, ) from lerobot.annotations.steerable_pipeline.reader import iter_episodes from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging from lerobot.annotations.steerable_pipeline.validator import StagingValidator from lerobot.annotations.steerable_pipeline.writer import LanguageColumnsWriter FALLBACK = ContextVar("original_prompt", default=False) CUSTOM_PROMPTS = ContextVar("custom_prompts", default=None) ORIGINAL_LOAD = plan_module.load_prompt ACTION_PROMPT = Path(__file__).with_name("action-phases.txt").read_text() def load_prompt(name): custom = CUSTOM_PROMPTS.get() or {} if name in custom: return custom[name] return ACTION_PROMPT if name == "plan_subtasks" and not FALLBACK.get() else ORIGINAL_LOAD(name) plan_module.load_prompt = load_prompt interjections_module.load_prompt = load_prompt vqa_module.load_prompt = load_prompt class FallbackClient: def __init__(self, client): self.client = client def generate_json(self, *args, **kwargs): kwargs["max_new_tokens"] = 512 return self.client.generate_json(*args, **kwargs) class TunedPlan(plan_module.PlanSubtasksMemoryModule): def _generate_plan(self, record, subtask_spans, *, refresh_t=None, interjection=None, task=None): custom = (CUSTOM_PROMPTS.get() or {}).get("plan_render") if not custom: return super()._generate_plan( record, subtask_spans, refresh_t=refresh_t, interjection=interjection, task=task ) remaining = [span for span in subtask_spans if refresh_t is None or float(span["start"]) >= refresh_t] if not remaining: return None text = custom.format( episode_task=task or record.episode_task, remaining_subtasks="\n".join(f"{i}. {s['text']}" for i, s in enumerate(remaining, 1)), ) result = self._vlm_field(self._text_message(text), "plan") if not isinstance(result, str) or not result.strip(): raise ValueError("Custom plan prompt did not return a nonempty plan string.") return result.strip() def _vlm_field(self, messages, field): response = self.vlm.generate_json([messages])[0] if isinstance(response, dict): return response.get(field) if ( field == "subtasks" and isinstance(response, list) and response and all( isinstance(s, dict) and isinstance(s.get("text"), str) and s["text"].strip() and all(type(s.get(k)) in (int, float) and math.isfinite(s[k]) for k in ("start", "end")) and s["end"] > s["start"] for s in response ) ): return response return None def _generate_subtasks(self, record, *, task=None): self.used_fallback = False try: result = super()._generate_subtasks(record, task=task) if not result: raise ValueError("No valid subtasks") return result except BudgetExceededError: raise except (ValueError, RuntimeError, KeyError, TypeError): self.used_fallback = True ctx = FALLBACK.set(True) try: fallback = plan_module.PlanSubtasksMemoryModule( vlm=FallbackClient(self.vlm), config=PlanConfig(), frame_provider=self.frame_provider ) result = fallback._generate_subtasks(record, task=task) if not result: raise ValueError("No valid subtasks after recovery") return result finally: FALLBACK.reset(ctx) class SourceFrames: """Exact timestamp seeks; a missing image fails the episode instead of generating from text.""" def __init__(self, row, info, root, camera, max_size=(224, 224)): self.row, self.info, self.root = row, info, root self.camera_key, self.camera_keys, self.cache = camera, [camera], {} # None keeps full resolution (the first frame of a human demo video). self.max_size = max_size def frames_at(self, record, timestamps, camera_key=None): camera = camera_key or self.camera_key if camera != self.camera_key: raise ValueError("Only the selected camera is available.") fps = float(self.info["fps"]) indices = [min(int(self.row["length"]) - 1, max(0, round(t * fps))) for t in timestamps] missing = sorted(set(indices) - self.cache.keys()) if missing: path = video_path(self.row, self.info, camera) offset = float(self.row[f"videos/{camera}/from_timestamp"]) with av.open(str(self.root / path)) as container: stream = container.streams.video[0] stream.codec_context.thread_count = 2 for idx in missing: target = offset + idx / fps container.seek(int(target / stream.time_base), stream=stream, backward=True) closest, distance = None, float("inf") for frame in container.decode(stream): if frame.pts is None: continue timestamp = float(frame.pts * stream.time_base) if abs(timestamp - target) < distance: closest, distance = frame, abs(timestamp - target) if timestamp >= target: break if closest is None or distance > 1 / fps + 1e-4: raise RuntimeError("Source video is missing an expected frame.") image = closest.to_image() if self.max_size: image.thumbnail(self.max_size) self.cache[idx] = image return [self.cache[i] for i in indices] def video_path(row, info, camera): return relative_path( info["video_path"].format( video_key=camera, chunk_index=int(row[f"videos/{camera}/chunk_index"]), file_index=int(row[f"videos/{camera}/file_index"]), ) ) class PreservingWriter(LanguageColumnsWriter): """Replace requested styles only, preserving every other episode and style.""" def __init__(self, episode_ids, styles): super().__init__(drop_existing_subtask_index=False) self.episode_ids, self.styles = set(episode_ids), set(styles) def _materialize_table(self, table, persistent, events, *, drop_old): ids = table.column("episode_index").to_pylist() for name, generated in (("language_persistent", persistent), ("language_events", events)): existing = table.column(name).to_pylist() if name in table.column_names else [[] for _ in ids] for i, ep in enumerate(ids): old = existing[i] or [] if ep not in self.episode_ids: generated[i] = old else: generated[i] = [r for r in old if r.get("style") not in self.styles] + generated[i] if name == "language_persistent": generated[i].sort(key=lambda r: (r.get("timestamp", 0), r.get("style") or "")) return super()._materialize_table(table, persistent, events, drop_old=drop_old) def requested_styles(features): # Human demo videos are media files, not language rows. styles = {FEATURES[f] for f in features if f != HUMAN_VIDEOS} if "Interjections & speech" in features: styles.add(None) return styles def annotate_episode(record, row, info, root, manifest, provider, staging_dir, media=None): context = CUSTOM_PROMPTS.set(validate_prompts(manifest.get("prompts", {}), manifest["features"])) try: return _annotate_episode(record, row, info, root, manifest, provider, staging_dir, media) finally: CUSTOM_PROMPTS.reset(context) def human_video_summary(entries): counts = defaultdict(int) for entry in entries: counts[entry["status"]] += 1 return dict(counts) def _annotate_episode(record, row, info, root, manifest, provider, staging_dir, media=None): frames = SourceFrames(row, info, root, manifest["camera"]) features = manifest["features"] staging = EpisodeStaging(staging_dir, record.episode_index) cfg = PlanConfig( emit_plan="Plans" in features or "Interjections & speech" in features, emit_memory="Memory" in features, n_task_rephrasings=10 if "Task rephrasings" in features else 0, max_frames_per_prompt=1200, plan_max_steps=80, min_subtask_seconds=0.5, subtask_describe_first=False, ) plan = TunedPlan(vlm=provider, config=cfg, frame_provider=frames) if set(features) - {"Visual questions"}: plan.run_episode(record, staging) if "Interjections & speech" in features: module = InterjectionsAndSpeechModule( vlm=provider, config=InterjectionsConfig(), frame_provider=frames ) module.run_episode(record, staging) if cfg.emit_plan: # Match Executor's plan-refresh pass while all dependency rows are still present. events = staging.read("interjections") refresh = [ (float(r["timestamp"]), r.get("content") or "") for r in events if r.get("style") == "interjection" ] plan.run_plan_updates(record, staging, [t for t, _ in refresh], [text for _, text in refresh]) if "Visual questions" in features: module = GeneralVqaModule( vlm=provider, config=VqaConfig(vqa_emission_hz=manifest["vqa_hz"], restrict_to_default_camera=True), frame_provider=frames, ) module.run_episode(record, staging) # Validate the complete dependency graph before removing unchecked output styles. report = StagingValidator(dataset_camera_keys=(manifest["camera"],)).validate([record], staging_dir) if not report.ok: raise ValueError("Generated annotations did not pass pipeline validation.") videos = None if HUMAN_VIDEOS in features and media is not None: # After validation, so no video is paid for on an episode whose annotations fail. full_frames = SourceFrames(row, info, root, manifest["camera"], max_size=None) try: videos = generate_human_videos( record, staging, provider, frames, full_frames, media, staging_dir.parent / "human_videos" / f"episode_{record.episode_index:06d}", manifest["human_video"]["segments_per_episode"], ) except BudgetExceededError as exc: # The annotations are complete; only the videos stop. status = "time_stopped" if isinstance(exc, DeadlineExceededError) else "budget_stopped" videos = [{"source_episode_index": record.episode_index, "status": status, "files": {}}] styles = requested_styles(features) for module_name, rows in staging.read_all().items(): staging.write(module_name, [r for r in rows if r.get("style") in styles]) result = { "episode_index": record.episode_index, "status": "completed", "recovery_used": getattr(plan, "used_fallback", False), "rows": sum(len(rows) for rows in staging.read_all().values()), } if videos is not None: result["human_videos"] = human_video_summary(videos) result["human_video_entries"] = videos return result HUMAN_VIDEO_DIR = "annotation_studio/human_videos" def video_deadline(deadline, manifest, now=None): """When videos must stop: the runtime limit if one was set, else the job timeout. Without a runtime limit the job still stops at its timeout (counted from job start, before setup), so keep a 15 minute margin for setup and the final commits. """ if deadline: return deadline now = time.time() if now is None else now return now + manifest["costs"]["hours"] * 3600 - 900 def human_video_operations(entries, index): """Upload an episode's media and the updated index (source episode IDs until materialized).""" operations = [] for entry in entries: files = entry.pop("files", {}) folder = f"{HUMAN_VIDEO_DIR}/episode_{entry['source_episode_index']:06d}" for kind, local in files.items(): name = f"{folder}/{Path(local).name}" operations.append(CommitOperationAdd(path_in_repo=name, path_or_fileobj=local)) entry[f"{kind}_path"] = name index.append(entry) operations.append(human_video_index_operation(index)) return operations def human_video_index_operation(index): return CommitOperationAdd( path_in_repo=f"{HUMAN_VIDEO_DIR}/index.json", path_or_fileobj=json.dumps(index, indent=2).encode() ) def create_output_repo(api, manifest): """Create an empty destination. Dataset copies are built only from completed episodes.""" visibility = manifest.get("visibility", "private") if visibility not in ("private", "public"): raise ValueError("Choose private or public output visibility.") if manifest.get("mode") != COPY: raise ValueError("Annotations must be saved in a LeRobot Parquet dataset.") output = manifest["output"] api.create_repo(output, repo_type="dataset", private=visibility == "private", exist_ok=False) license_id = manifest.get("source_license") license_line = f"license: {json.dumps(license_id)}\n" if isinstance(license_id, str) else "" kind = "Annotated robot episodes" description = ( "Only successfully annotated episodes are included, with all RGB camera videos trimmed to those " "episodes. Episode and frame indices are rebuilt. See `annotation_studio/episode_mapping.json` " "for the original episode IDs. RGB videos are re-encoded as H.264 (CRF 18)." ) card = ( f"---\nbase_datasets:\n- {manifest['dataset']}\n{license_line}tags:\n- lerobot\n- annotation\n---\n\n# {kind}\n\n" f"Annotations for [{manifest['dataset']}](https://huggingface.co/datasets/{manifest['dataset']}) " f"at `{manifest['dataset_revision']}`. {description} " "See `annotation_studio/run.json` for completeness.\n" ) if HUMAN_VIDEOS in manifest.get("features", ()): profile = manifest["human_video"]["profile"] card += ( "\n## Human demonstration videos\n\n" "`annotation_studio/human_videos/` holds a generated first-person human video for " + ( f"up to {manifest['human_video']['segments_per_episode']} subtasks" if manifest["human_video"]["segments_per_episode"] else "every subtask" ) + " per episode: the robot frame, " + f"the {profile['edit_model']} edit and a {profile['resolution']} {profile['video_model']} " "video. `index.json` lists the subtask, timestamps, prompts and status of every segment. " "These videos are synthetic.\n" ) api.upload_file( path_or_fileobj=card.encode(), path_in_repo="README.md", repo_id=output, repo_type="dataset" ) def run(manifest, work=Path("/workspace/run")): token = os.environ["HF_TOKEN"] api = HfApi(token=token) work.mkdir(parents=True, exist_ok=True) selection = read_selection( manifest["dataset"], manifest["start"], manifest["count"], token, revision=manifest["dataset_revision"], cache_dir=str(work / "metadata-cache"), ) ids = [r["episode_index"] for r in selection["episodes"]] if hashlib.sha256(json.dumps(ids).encode()).hexdigest() != manifest["episode_ids_sha256"]: raise ValueError("Pinned episode selection differs from the estimate.") output = manifest["output"] create_output_repo(api, manifest) root, staging = work / "dataset", work / "staging" snapshot_download( manifest["dataset"], repo_type="dataset", revision=manifest["dataset_revision"], token=token, local_dir=root, allow_patterns=["meta/**", "README.md", "LICENSE*", "COPYING*"], ) subset = EpisodeSubset(root, work / "subset", ids) attribution = [] for file in root.iterdir(): if file.is_file() and (file.name == "README.md" or file.name.startswith(("LICENSE", "COPYING"))): target = subset.root / ( "annotation_studio/source_README.md" if file.name == "README.md" else file.name ) target.parent.mkdir(parents=True, exist_ok=True) shutil.copyfile(file, target) attribution.append(target) budget = Budget( manifest["inference_budget"], manifest["costs"]["input_rate"], manifest["costs"]["output_rate"], manifest["costs"].get("cached_input_rate"), context_length=manifest["costs"].get("context_length", 262144), pricing_tiers=manifest["costs"].get("pricing_tiers"), ) deadline = float(os.environ.get("ANNOTATION_DEADLINE_EPOCH", "0")) or None provider = Provider( token, manifest["model"], budget, deadline=deadline, concurrency=manifest.get("request_concurrency", 2), generation_options=manifest.get("generation_options"), ) media = ( FalMedia( token, budget, video_deadline(deadline, manifest), manifest["human_video"]["profile"], concurrency=manifest["human_video"].get("parallel"), ) if HUMAN_VIDEOS in manifest["features"] and manifest.get("human_video") else None ) human_index = [] progress = { "manifest": manifest, "pipeline": json.loads(Path(__file__).with_name("pipeline-version.json").read_text()), "status": "running", "episodes": [], "expected_episodes": len(ids), "written_episode_ids": [], "source_episode_map": [], "copy_scope": "annotated_episodes", "worker_started_at": time.time(), "request_concurrency": provider.concurrency, } commit_lock = threading.RLock() def checkpoint(extra=()): with commit_lock: return _checkpoint(extra) def _checkpoint(extra=()): progress["usage"] = budget.summary() operations = [ CommitOperationAdd( path_in_repo="annotation_studio/run.json", path_or_fileobj=json.dumps(progress, indent=2).encode(), ), *extra, ] return api.create_commit( output, repo_type="dataset", operations=operations, commit_message=f"Annotation progress: {len(progress['episodes'])}/{len(ids)}", ) checkpoint() heartbeat_stop = threading.Event() def heartbeat(): while not heartbeat_stop.wait(10): try: checkpoint() except Exception: print("Progress update delayed; will retry.", flush=True) heartbeat_thread = threading.Thread(target=heartbeat, daemon=True) heartbeat_thread.start() shards = defaultdict(list) info = selection["info"] for row in selection["episodes"]: path = relative_path( info["data_path"].format( chunk_index=int(row["data/chunk_index"]), file_index=int(row["data/file_index"]) ) ) shards[path].append(row) try: for data_path, rows in shards.items(): files = {data_path} | {video_path(row, info, manifest["camera"]) for row in rows} sizes = api.get_paths_info( manifest["dataset"], list(files), repo_type="dataset", revision=manifest["dataset_revision"] ) if sum(x.size for x in sizes) * 2 + 2_000_000_000 > shutil.disk_usage(work).free: raise RuntimeError("Selected shard is too large for this CPU Job's disk.") for file in sorted(files): hf_hub_download( manifest["dataset"], file, repo_type="dataset", revision=manifest["dataset_revision"], token=token, local_dir=root, ) row_by_id = {int(r["episode_index"]): r for r in rows} records = list(iter_episodes(root, only_episodes=tuple(row_by_id))) if {r.episode_index for r in records} != set(row_by_id): raise ValueError("Data shards do not match episode metadata.") completed = [] stop = False parallelism = max(1, min(manifest.get("episode_parallelism", 2), provider.concurrency)) with ThreadPoolExecutor(max_workers=parallelism) as pool: futures = { pool.submit( annotate_episode, record, row_by_id[record.episode_index], info, root, manifest, provider, staging, media, ): record for record in records } for future in as_completed(futures): record = futures[future] if future.cancelled(): continue try: result = future.result() completed.append(record) except BudgetExceededError as exc: result = { "episode_index": record.episode_index, "status": "time_stopped" if isinstance(exc, DeadlineExceededError) else "budget_stopped", } stop = True for pending in futures: pending.cancel() except Exception as exc: result = { "episode_index": record.episode_index, "status": "failed", "error_type": type(exc).__name__, } entries = result.pop("human_video_entries", None) progress["episodes"].append(result) if entries: local_files = [f for e in entries for f in e.get("files", {}).values()] checkpoint(human_video_operations(entries, human_index)) for local in local_files: Path(local).unlink(missing_ok=True) else: checkpoint() print(f"Saved episode {record.episode_index}: {result['status']}", flush=True) if completed: writer = PreservingWriter( [r.episode_index for r in completed], requested_styles(manifest["features"]) ) writer.write_all(completed, staging, root) Executor._ensure_annotation_metadata_in_info(root) annotated_table = pq.read_table(root / data_path) for record in sorted(completed, key=lambda r: r.episode_index): row = row_by_id[record.episode_index] # Other cameras are needed only for episodes whose annotations succeeded. extra_files = {video_path(row, info, camera) for camera in subset.cameras} - files needed = [name for name in extra_files if not (root / name).exists()] sizes = ( api.get_paths_info( manifest["dataset"], needed, repo_type="dataset", revision=manifest["dataset_revision"], ) if needed else [] ) if sum(item.size for item in sizes) * 2 + 2_000_000_000 > shutil.disk_usage(work).free: raise RuntimeError("Episode camera files are too large for this CPU Job's disk.") for name in sorted(needed): hf_hub_download( manifest["dataset"], name, repo_type="dataset", revision=manifest["dataset_revision"], token=token, local_dir=root, ) new_id, written = subset.add(record.episode_index, annotated_table) extra_operations = [] for entry in human_index: if entry["source_episode_index"] == record.episode_index: entry["episode_index"] = new_id extra_operations = [human_video_index_operation(human_index)] with commit_lock: previous_ids = list(progress["written_episode_ids"]) previous_map = list(progress["source_episode_map"]) progress["written_episode_ids"].append(new_id) progress["source_episode_map"] = list(subset.mapping) try: checkpoint( [ CommitOperationAdd( path_in_repo=str(p.relative_to(subset.root)), path_or_fileobj=p ) for p in [*written, *attribution] ] + extra_operations ) except Exception: progress["written_episode_ids"] = previous_ids progress["source_episode_map"] = previous_map raise attribution = [] # Release clips immediately; published episodes no longer consume local disk. for path in written: if path.suffix == ".mp4" or ("data" in path.relative_to(subset.root).parts): path.unlink(missing_ok=True) for name in extra_files: (root / name).unlink(missing_ok=True) for file in files: (root / file).unlink(missing_ok=True) shutil.rmtree(staging, ignore_errors=True) if stop: break heartbeat_stop.set() heartbeat_thread.join(timeout=15) progress["worker_finished_at"] = time.time() completed_count = sum(e["status"] == "completed" for e in progress["episodes"]) progress["status"] = "completed" if completed_count == len(ids) else "partial" final_commit = checkpoint() if progress["written_episode_ids"]: # Tag only a usable materialized subset, including a partial successful run. version = str(info["codebase_version"]) tags = api.list_repo_refs(output, repo_type="dataset").tags if any(t.name == version for t in tags): api.delete_tag(output, tag=version, repo_type="dataset") api.create_tag(output, tag=version, repo_type="dataset", revision=final_commit.oid) except Exception as exc: heartbeat_stop.set() heartbeat_thread.join(timeout=15) progress.update(status="partial", error_type=type(exc).__name__, worker_finished_at=time.time()) checkpoint() raise RuntimeError("Annotation stopped; check the saved run report for completed episodes.") from None print(f"Result: {progress['status']} — https://huggingface.co/datasets/{output}", flush=True) if __name__ == "__main__": run(json.loads(os.environ["ANNOTATION_MANIFEST"]))