| """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], {} |
| |
| 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): |
| |
| 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: |
| |
| 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) |
| |
| 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: |
| |
| 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: |
| |
| 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] |
| |
| 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 = [] |
| |
| 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"]: |
| |
| 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"])) |
|
|
