| """Studio episode/media stages on the LeRobot #4892 execution contracts. |
| |
| The CPU controller keeps its shared billing guard and provider semaphores in |
| process. Only opaque in-process handles enter operational plans; tokens and |
| provider objects are never serialized. Scientific identities exclude handles. |
| """ |
|
|
| import hashlib |
| import json |
| import shutil |
| import tarfile |
| import uuid |
| from dataclasses import asdict |
| from pathlib import Path |
|
|
| from lerobot.annotations.runtime.configs import RuntimeConfig, StageConfig |
| from lerobot.annotations.runtime.pipeline import run_pipeline |
| from lerobot.annotations.runtime.types import ( |
| Artifact, |
| DatasetRef, |
| InputItem, |
| ItemResult, |
| ModuleSpec, |
| Outcome, |
| fingerprint, |
| ) |
| from lerobot.annotations.runtime.worker import accepted_items |
|
|
| CONTEXTS = {} |
|
|
|
|
| def unpack(bundle, directory): |
| with tarfile.open(bundle) as archive: |
| for member in archive.getmembers(): |
| path = Path(member.name) |
| if not member.isfile() or path.is_absolute() or ".." in path.parts: |
| raise ValueError("Invalid Studio checkpoint bundle") |
| target = directory / path |
| target.parent.mkdir(parents=True, exist_ok=True) |
| target.write_bytes(archive.extractfile(member).read()) |
|
|
|
|
| class StudioStage: |
| def __init__(self, handle, phase, scientific): |
| self.handle, self.phase, self.scientific = handle, phase, scientific |
| self.spec = ModuleSpec("studio_" + phase, "1", "episode", {"bundle": None}) |
|
|
| def scientific_config(self): |
| return {"phase": self.phase, **self.scientific} |
|
|
| def setup(self, context): |
| self.values = CONTEXTS[self.handle] |
|
|
| def teardown(self): |
| pass |
|
|
| def process_batch(self, items, context): |
| import worker |
|
|
| from lerobot.annotations.steerable_pipeline.staging import EpisodeStaging |
|
|
| record, row, info, root, manifest, provider, media = self.values |
| results = [] |
| for item in items: |
| directory = context.scratch / item.item_id |
| directory.mkdir(parents=True, exist_ok=True) |
| staging_dir = directory / "staging" |
| if self.phase == "language": |
| result = worker._annotate_episode( |
| record, row, info, root, manifest, provider, staging_dir, None |
| ) |
| else: |
| artifact = Artifact(**item.payload["upstream"]) |
| if not context.store.verify(artifact): |
| raise ValueError("Upstream checkpoint checksum mismatch") |
| bundle = context.store.download(artifact, directory / "language.tar.gz") |
| unpack(bundle, directory) |
| bundle.unlink() |
| result = json.loads((directory / "result.json").read_text()) |
| if self.phase == "hands": |
| import math |
| import time |
|
|
| import av |
| import pyarrow.parquet as pq |
| from subset import trim_video |
|
|
| from lerobot.annotations.hands.writer import attach_actions |
|
|
| if worker.HAND_ANNOTATOR is None: |
| raise RuntimeError("The HF Jobs hand worker was not initialized") |
| native = directory / "hand-native" |
| native.mkdir(parents=True, exist_ok=True) |
| video = native / "episode.mp4" |
| started = time.perf_counter() |
| print( |
| "HAND_RUNTIME " |
| + json.dumps( |
| {"event": "episode_video_prepare_start", "source_episode": record.episode_index} |
| ), |
| flush=True, |
| ) |
| source_video = root / worker.video_path(row, info, manifest["camera"]) |
| offset = float(row[f"videos/{manifest['camera']}/from_timestamp"]) |
| |
| |
| with av.open(str(source_video)) as probe: |
| stream = probe.streams.video[0] |
| complete_episode = ( |
| len(probe.streams.video) == 1 |
| and offset == 0 |
| and stream.start_time == 0 |
| and stream.frames == int(row["length"]) |
| and math.isclose(float(stream.average_rate or 0), float(info["fps"]), rel_tol=1e-6) |
| ) |
| if complete_episode: |
| video = source_video |
| else: |
| trim_video( |
| source_video, video, offset, int(row["length"]), float(info["fps"]), lossless=True |
| ) |
| prepared = time.perf_counter() |
| print( |
| "HAND_RUNTIME " |
| + json.dumps( |
| { |
| "event": "native_annotation_start", |
| "source_episode": record.episode_index, |
| "video_prepare_seconds": prepared - started, |
| "input_mode": "whole-episode-video" if complete_episode else "lossless-trim", |
| } |
| ), |
| flush=True, |
| ) |
| try: |
| native_report = worker.HAND_ANNOTATOR.annotate( |
| video, |
| manifest["hand_actions"]["calibration"], |
| native / "work", |
| orientation=manifest["hand_actions"].get("orientation", "sensor"), |
| ) |
| except Exception: |
| |
| import logging |
|
|
| logging.exception( |
| "Native hand annotation failed for source episode %s", record.episode_index |
| ) |
| raise |
| inferred = time.perf_counter() |
| table = pq.read_table( |
| record.data_path, filters=[("episode_index", "=", record.episode_index)] |
| ) |
| table, summary = attach_actions(table, native_report, directory=native, fps=info["fps"]) |
| summary["timing_seconds"] = { |
| "source_video_prepare": prepared - started, |
| "native_annotation": inferred - prepared, |
| "native_scope": "intrinsics, rectification, camera, hands and native export; includes coordinator wait and first-use model initialization; excludes source preparation and final dataset video copying", |
| } |
| summary["native_components"] = { |
| key: native_report[key] |
| for key in ("camera_statistics", "hand_components") |
| if key in native_report |
| } |
| print( |
| "HAND_RUNTIME " |
| + json.dumps( |
| { |
| "event": "native_annotation_finished", |
| "source_episode": record.episode_index, |
| "native_seconds": inferred - prepared, |
| "valid_frames": summary["valid_frames"], |
| } |
| ), |
| flush=True, |
| ) |
| pq.write_table(table, directory / "hand_actions.parquet") |
| result["hand_actions"] = summary |
| |
| shutil.rmtree(native) |
| elif self.phase == "human_video": |
| staging = EpisodeStaging(staging_dir, record.episode_index) |
| frames = worker.SourceFrames(row, info, root, manifest["camera"]) |
| full = worker.SourceFrames(row, info, root, manifest["camera"], max_size=None) |
| try: |
| entries = worker.generate_human_videos( |
| record, |
| staging, |
| provider, |
| frames, |
| full, |
| media, |
| directory / "human_videos" / f"episode_{record.episode_index:06d}", |
| manifest["human_video"]["segments_per_episode"], |
| ) |
| except worker.BudgetExceededError as exc: |
| entries = [ |
| { |
| "source_episode_index": record.episode_index, |
| "status": "time_stopped" |
| if isinstance(exc, worker.DeadlineExceededError) |
| else "budget_stopped", |
| "files": {}, |
| } |
| ] |
| for entry in entries: |
| entry["files"] = { |
| kind: str(Path(path).relative_to(directory)) |
| for kind, path in entry.get("files", {}).items() |
| } |
| result["human_videos"] = worker.human_video_summary(entries) |
| result["human_video_entries"] = entries |
| (directory / "result.json").write_text(json.dumps(result)) |
| bundle = context.scratch / (item.item_id + ".tar.gz") |
| with tarfile.open(bundle, "w:gz") as archive: |
| for path in sorted(directory.rglob("*")): |
| if path.is_file(): |
| archive.add(path, arcname=str(path.relative_to(directory))) |
| artifact = context.write_asset(item, "bundle", bundle) |
| results.append(ItemResult(item.item_id, Outcome.COMPLETED, (artifact,))) |
| return results |
|
|
|
|
| class EpisodeSource: |
| def __init__(self, record, row, info, manifest): |
| self.record, self.row = record, row |
| self.identity = { |
| "row": row, |
| "fps": info["fps"], |
| "task": record.episode_task, |
| "source": manifest.get("dataset"), |
| "revision": manifest.get("dataset_revision"), |
| } |
| revision = manifest.get("dataset_revision") or fingerprint(self.identity) |
| self.dataset_ref = DatasetRef(manifest.get("dataset", "studio-local"), revision) |
|
|
| def discover(self, stage, store, upstream): |
| payload = {} |
| identity = dict(self.identity) |
| if upstream: |
| plan, _ = next(iter(upstream.values())) |
| _, result = next(accepted_items(store, plan)) |
| artifact = result.artifacts[0] |
| payload["upstream"] = asdict(artifact) |
| identity["upstream"] = {key: value for key, value in asdict(artifact).items() if key != "path"} |
| yield InputItem(str(self.record.episode_index), payload, content_digest=fingerprint(identity)) |
|
|
|
|
| def run_episode(record, row, info, root, manifest, provider, staging_dir, media=None): |
| handle = uuid.uuid4().hex |
| CONTEXTS[handle] = (record, row, info, root, manifest, provider, media) |
| folder = Path(__file__).parent |
| scientific = { |
| "options": { |
| key: manifest.get(key) |
| for key in ( |
| "features", |
| "camera", |
| "vqa_hz", |
| "prompts", |
| "model", |
| "generation_options", |
| "human_video", |
| "captioner", |
| "captioner_profile", |
| "flora_execution", |
| "hand_actions", |
| ) |
| }, |
| "code": { |
| path.name: hashlib.sha256(path.read_bytes()).hexdigest() |
| for pattern in ("*.py", "*.txt", "pipeline-version.json") |
| for path in sorted(folder.glob(pattern)) |
| if not path.name.startswith("test_") |
| }, |
| } |
| stages = [ |
| StageConfig( |
| "language", |
| "studio_runtime:StudioStage", |
| {"handle": handle, "phase": "language", "scientific": scientific}, |
| ) |
| ] |
| if media is not None: |
| stages.append( |
| StageConfig( |
| "human_video", |
| "studio_runtime:StudioStage", |
| {"handle": handle, "phase": "human_video", "scientific": scientific}, |
| ("language",), |
| ) |
| ) |
| if manifest.get("hand_actions"): |
| stages.append( |
| StageConfig( |
| "hands", |
| "studio_runtime:StudioStage", |
| {"handle": handle, "phase": "hands", "scientific": scientific}, |
| (stages[-1].id,), |
| ) |
| ) |
| |
| |
| runtime = RuntimeConfig( |
| run_uri=str(staging_dir.parent / "runtime" / f"episode_{record.episode_index:06d}"), |
| workers=1, |
| batch_size=1, |
| shard_size=1, |
| max_retries=0, |
| ) |
| try: |
| store, completed = run_pipeline(EpisodeSource(record, row, info, manifest), stages, runtime) |
| plan, _ = completed[stages[-1].id] |
| _, result = next(accepted_items(store, plan)) |
| directory = staging_dir.parent / "restored" / handle |
| directory.mkdir(parents=True, exist_ok=True) |
| bundle = store.download(result.artifacts[0], directory / "result.tar.gz") |
| unpack(bundle, directory) |
| result = json.loads((directory / "result.json").read_text()) |
| if (directory / "hand_actions.parquet").exists(): |
| target = staging_dir.parent / "hand_actions" / f"episode_{record.episode_index:06d}.parquet" |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(directory / "hand_actions.parquet", target) |
| if (directory / "staging").exists(): |
| shutil.rmtree(staging_dir / f"episode_{record.episode_index:06d}", ignore_errors=True) |
| shutil.copytree(directory / "staging", staging_dir, dirs_exist_ok=True) |
| for entry in result.get("human_video_entries", []): |
| for kind, relative in entry.get("files", {}).items(): |
| source = directory / relative |
| target = staging_dir.parent / relative |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(source, target) |
| entry["files"][kind] = str(target) |
| result["runtime"] = {name: summary.plan_id for name, (_, summary) in completed.items()} |
| shutil.rmtree(directory) |
| return result |
| finally: |
| CONTEXTS.pop(handle, None) |
|
|
