Spaces:

finerobotics
/

annotation-studio

Running

Preserve bounded hand overlap and reliable native Hub downloads
ed5344e verified 6 days ago
43 kB
"""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 sys
import tempfile
import threading
import time
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from contextvars import ContextVar
from dataclasses import replace
from pathlib import Path

import av
import pyarrow.parquet as pq
from hand_actions import HANDS
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, needs_provider, read_selection, relative_path
from studio_runtime import run_episode as run_runtime_episode
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)
STUDIO_ACTIVE = ContextVar("studio_active", default=False)
ORIGINAL_LOAD = plan_module.load_prompt
ACTION_PROMPT = Path(__file__).with_name("action-phases.txt").read_text()
FLORA_CAPTIONER = None
HAND_ANNOTATOR = None


class IndependentFloraClient:
    """Fail closed if a subtask-only Flora workflow attempts a provider request."""

    concurrency = 1

    def generate_json(self, *args, **kwargs):
        raise RuntimeError("Independent Flora subtasks must not call an external VLM")


def load_prompt(name):
    custom = CUSTOM_PROMPTS.get() or {}
    if not STUDIO_ACTIVE.get() and not custom:
        return ORIGINAL_LOAD(name)
    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_subtasks(self, record, *, task=None):
        context = STUDIO_ACTIVE.set(True)
        try:
            if getattr(self, "captioner", None) is not None:
                self.used_fallback = False
                spans = self.captioner.annotate(record, self.caption_frames, self.caption_fps)
            else:
                spans = self._segment_subtasks(record, task=task)
            self.subtask_spans = spans
            return spans
        finally:
            STUDIO_ACTIVE.reset(context)

    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 _segment_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
                decoded, current, previous = None, None, None
                for idx in missing:
                    target = offset + idx / fps
                    # Native Flora requests dense, sorted samples. Reuse the
                    # decoder instead of decoding the same GOP for each frame.
                    # Seek again across sparse gaps so those remain cheap too.
                    if current is None or target - float(current.pts * stream.time_base) > 1:
                        container.seek(int(target / stream.time_base), stream=stream, backward=True)
                        decoded, current, previous = iter(container.decode(stream)), None, None
                    while current is None or float(current.pts * stream.time_base) < target:
                        frame = next(decoded, None)
                        if frame is None:
                            break
                        if frame.pts is not None:
                            previous, current = current, frame
                    candidates = [frame for frame in (previous, current) if frame is not None]
                    closest = min(
                        candidates,
                        key=lambda frame: abs(float(frame.pts * stream.time_base) - target),
                        default=None,
                    )
                    distance = (
                        abs(float(closest.pts * stream.time_base) - target) if closest else float("inf")
                    )
                    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"]),
            episode_index=row.get("episode_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 not in {HUMAN_VIDEOS, HANDS}}
    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"]))
    active = STUDIO_ACTIVE.set(True)
    try:
        return run_runtime_episode(record, row, info, root, manifest, provider, staging_dir, media)
    finally:
        STUDIO_ACTIVE.reset(active)
        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 manifest.get("captioner") == "flora":
        if FLORA_CAPTIONER is None:
            raise RuntimeError("Flora captioner is not initialized on this worker")
        plan.captioner = FLORA_CAPTIONER
        # Even an empty/placeholder task must not trigger a Qwen call before Flora.
        cfg.derive_task_from_video = "off"
        plan.caption_frames = SourceFrames(row, info, root, manifest["camera"], max_size=None)
        plan.caption_fps = info["fps"]
    if set(features) - {"Visual questions", HANDS}:
        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()),
        # Parquet language rows store starts only. Keep native ends, gaps and
        # overlaps in the run report so the two independent runs can be compared.
        "subtask_spans": [
            {
                **{key: span[key] for key in ("start", "end", "text")},
                **(
                    {"segmentation_provenance": span["metadata"]["segmentation_provenance"]}
                    if span.get("metadata", {}).get("segmentation_provenance")
                    else {}
                ),
            }
            for span in getattr(plan, "subtask_spans", [])
        ],
    }
    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 prefetch_source_metadata(manifest, work, token):
    """Overlap read-only source setup with native model downloads and verification."""
    started = time.monotonic()
    selection = read_selection(
        manifest["dataset"],
        manifest["start"],
        manifest["count"],
        token,
        revision=manifest["dataset_revision"],
        cache_dir=str(work / "metadata-cache"),
    )
    snapshot_download(
        manifest["dataset"],
        repo_type="dataset",
        revision=manifest["dataset_revision"],
        token=token,
        local_dir=work / "dataset",
        allow_patterns=["meta/**", "README.md", "LICENSE*", "COPYING*"],
    )
    print(
        "FLORA_RUNTIME "
        + json.dumps({"event": "source_metadata_ready", "elapsed_s": round(time.monotonic() - started, 3)}),
        flush=True,
    )
    return selection


def run(manifest, work=Path("/workspace/run")):
    global FLORA_CAPTIONER, HAND_ANNOTATOR
    if (HANDS in manifest["features"]) != bool(manifest.get("hand_actions")):
        raise ValueError("Hand action selection and pinned runtime configuration differ")
    deadline = float(os.environ.get("ANNOTATION_DEADLINE_EPOCH", "0")) or None
    source_pool = source_future = None
    if manifest.get("captioner") == "flora":
        from flora_captioner import FLORA, FloraCaptioner

        if manifest.get("captioner_profile") != FLORA:
            raise ValueError("Unrecognized Flora revision or configuration")
        FLORA_CAPTIONER = FloraCaptioner(
            deadline=deadline,
            cache_models=FLORA["resident_models"],
            execution=os.environ.get("FLORA_EXECUTION", "reference"),
            batch_size=int(os.environ.get("FLORA_BATCH_SIZE", "4")),
            vision_frame_batching=os.environ.get("FLORA_VISION_FRAME_BATCHING", "0") == "1",
        )
        flora_parallelism = int(os.environ.get("FLORA_EPISODE_PARALLELISM", "1"))
        if flora_parallelism > 1:
            if FLORA_CAPTIONER.execution != "accelerated" or manifest["features"] != ["Subtasks"]:
                raise ValueError("Cross-episode Flora requires independent accelerated subtasks only")
            from flora_batcher import CrossEpisodeFlora

            work.mkdir(parents=True, exist_ok=True)
            source_pool = ThreadPoolExecutor(max_workers=1)
            source_future = source_pool.submit(
                prefetch_source_metadata, manifest, work, os.environ["HF_TOKEN"]
            )
            try:
                FLORA_CAPTIONER = CrossEpisodeFlora(
                    FLORA_CAPTIONER,
                    gpu_batch_size=int(os.environ.get("FLORA_GPU_BATCH_SIZE", "64")),
                    parallelism=flora_parallelism,
                )
            except BaseException:
                source_pool.shutdown(wait=True, cancel_futures=True)
                raise
        manifest["flora_execution"] = {
            "profile": FLORA_CAPTIONER.execution,
            "batch_size": FLORA_CAPTIONER.batch_size,
            "kernel_versions": FLORA_CAPTIONER.kernel_versions,
            "episode_parallelism": flora_parallelism,
            "generation_batch_size": getattr(FLORA_CAPTIONER, "gpu_batch_size", FLORA_CAPTIONER.batch_size),
            "vision_frame_batching": getattr(FLORA_CAPTIONER, "vision_frame_batching", False),
            "shape_padding": "minimum_1792_tempseg64_captioner16"
            if getattr(FLORA_CAPTIONER, "vision_frame_batching", False)
            else "none",
        }
    token = os.environ["HF_TOKEN"]
    api = HfApi(token=token)
    work.mkdir(parents=True, exist_ok=True)
    if source_future is not None:
        try:
            selection = source_future.result()
        finally:
            source_pool.shutdown(wait=True)
    else:
        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"
    if source_future is None:
        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, camera_workers=3 if manifest.get("captioner") == "flora" else 1
    )
    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"),
    )
    provider = (
        Provider(
            token,
            manifest["model"],
            budget,
            deadline=deadline,
            concurrency=manifest.get("request_concurrency", 2),
            generation_options=manifest.get("generation_options"),
        )
        if needs_provider(manifest["features"], manifest.get("prompts"), manifest.get("captioner"))
        else IndependentFloraClient()
    )
    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:
        if manifest.get("hand_actions"):
            from hand_setup import prepare_camera_engines, verify_image

            from lerobot.annotations.hands.assets import download_release, verify_assets
            from lerobot.annotations.hands.config import HandBackendConfig
            from lerobot.annotations.hands.processing import HandAnnotator

            hands = manifest["hand_actions"]
            # Independent hands have no hosted requests to throttle. Keep the
            # native coordinator's bounded two-clip overlap available.
            if set(manifest["features"]) == {HANDS}:
                provider.concurrency = 2
            verify_image(hands["runtime"])
            backend = HandBackendConfig(**hands["runtime"]["backend"])
            print(
                "HAND_RUNTIME "
                + json.dumps(
                    {"event": "asset_download_start", "revision": hands["runtime"]["assets"]["revision"]}
                ),
                flush=True,
            )
            hands["asset_checksums"] = download_release(
                hands["runtime"]["assets"], token, Path(backend.assets_manifest).parent
            )
            verify_assets(backend)
            print(
                "HAND_RUNTIME " + json.dumps({"event": "assets_verified", **hands["asset_checksums"]}),
                flush=True,
            )
            prepare_camera_engines(hands["runtime"])
            print('HAND_RUNTIME {"event":"camera_engines_ready","native_worker":"starting"}', flush=True)
            HAND_ANNOTATOR = HandAnnotator(backend, work / "hands-native")
        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.")
            Executor._ensure_annotation_metadata_in_info(root)
            export_errors = []

            def export_episode(
                record, *, shard_path=data_path, shard_rows=row_by_id, shard_files=files, errors=export_errors
            ):
                nonlocal attribution
                if errors:
                    return
                try:
                    # Write a private episode slice. Inference threads keep
                    # reading the immutable source shard while this is exported.
                    with tempfile.TemporaryDirectory(dir=work, prefix="export-") as temporary:
                        table = pq.read_table(
                            root / shard_path, filters=[("episode_index", "=", record.episode_index)]
                        )
                        prepared_path = Path(temporary) / "episode.parquet"
                        pq.write_table(table, prepared_path)
                        prepared = replace(record, data_path=prepared_path, data_paths=(), row_offset=0)
                        writer = PreservingWriter(
                            [record.episode_index], requested_styles(manifest["features"])
                        )
                        writer.write_all([prepared], staging, root)
                        annotated_table = pq.read_table(prepared_path)
                    row = shard_rows[record.episode_index]
                    extra_files = {video_path(row, info, camera) for camera in subset.cameras} - shard_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 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,
                        )
                    shard_files.update(extra_files)
                    hand_features = None
                    if manifest.get("hand_actions"):
                        from lerobot.annotations.hands.writer import FEATURES as HAND_FEATURES

                        hand_features = HAND_FEATURES

                        hand_path = (
                            staging.parent / "hand_actions" / f"episode_{record.episode_index:06d}.parquet"
                        )
                        hand_table = pq.read_table(hand_path)
                        if not hand_table["frame_index"].equals(annotated_table["frame_index"]):
                            raise ValueError("Hand actions do not match the exported episode")
                        for name in [*HAND_FEATURES, "annotation.original_action"]:
                            if name in hand_table.column_names:
                                index = annotated_table.schema.get_field_index(name)
                                annotated_table = (
                                    annotated_table.set_column(index, name, hand_table[name])
                                    if index >= 0
                                    else annotated_table.append_column(name, hand_table[name])
                                )
                    new_id, written = subset.add(
                        record.episode_index, annotated_table, hand_features=hand_features
                    )
                    extra_operations = []
                    with commit_lock:
                        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)]
                        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(path.relative_to(subset.root)), path_or_fileobj=path
                                    )
                                    for path in [*written, *attribution]
                                ]
                                + extra_operations
                            )
                        except Exception:
                            progress["written_episode_ids"] = previous_ids
                            progress["source_episode_map"] = previous_map
                            raise
                    attribution = []
                    for path in written:
                        if path.suffix == ".mp4" or "data" in path.relative_to(subset.root).parts:
                            path.unlink(missing_ok=True)
                    print(f"Exported episode {record.episode_index} as {new_id}", flush=True)
                except Exception as exc:
                    errors.append(exc)
                    progress["export_error"] = {
                        "source_episode": record.episode_index,
                        "type": type(exc).__name__,
                    }

            # A single publisher owns subset indices and shared metadata. Its
            # CPU video work overlaps GPU inference without mutating the source.
            processed, next_export, submitted_exports = {}, 0, set()
            stream_completed_flora = manifest.get("captioner") == "flora" and hasattr(
                FLORA_CAPTIONER, "parallelism"
            )

            def export_ready(
                publisher,
                *,
                batch_records=records,
                statuses=processed,
                export=export_episode,
                unordered=stream_completed_flora,
                submitted=submitted_exports,
            ):
                nonlocal next_export
                if unordered:
                    # An unfinished early episode must not block video export
                    # for the rest. EpisodeSubset records the original IDs.
                    for record in batch_records:
                        source_id = record.episode_index
                        if statuses.get(source_id) == "completed" and source_id not in submitted:
                            publisher.submit(export, record)
                            submitted.add(source_id)
                    return
                while (
                    next_export < len(batch_records) and batch_records[next_export].episode_index in statuses
                ):
                    record = batch_records[next_export]
                    if statuses[record.episode_index] == "completed":
                        publisher.submit(export, record)
                    next_export += 1

            stop = False
            parallelism = (
                FLORA_CAPTIONER.parallelism
                if manifest.get("captioner") == "flora" and hasattr(FLORA_CAPTIONER, "parallelism")
                else max(1, min(manifest.get("episode_parallelism", 2), provider.concurrency))
            )
            with (
                ThreadPoolExecutor(max_workers=1) as publisher,
                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():
                        processed[record.episode_index] = "canceled"
                        export_ready(publisher)
                        continue
                    try:
                        result = future.result()
                    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)
                    processed[record.episode_index] = result["status"]
                    export_ready(publisher)
            if export_errors:
                raise export_errors[0]
            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
    finally:
        if HAND_ANNOTATOR is not None:
            HAND_ANNOTATOR.close()
            HAND_ANNOTATOR = None
    print(f"Result: {progress['status']} — https://huggingface.co/datasets/{output}", flush=True)


if __name__ == "__main__":
    # Runtime factories import worker by name and must share its prompt context.
    sys.modules["worker"] = sys.modules[__name__]
    run(json.loads(os.environ["ANNOTATION_MANIFEST"]))