from fastapi import FastAPI from fastapi import HTTPException from fastapi.middleware.cors import CORSMiddleware from fastapi.responses import StreamingResponse from fastapi.staticfiles import StaticFiles from huggingface_hub import HfApi from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_url from huggingface_hub.utils import EntryNotFoundError from huggingface_hub.utils import RepositoryNotFoundError from pydantic import BaseModel from typing import Optional import csv from datetime import datetime from datetime import timezone import io import json import os from pathlib import Path import random import re import threading import time from uuid import uuid4 # import assignment from backend import assignment from backend.validate_prompts import check_items app = FastAPI() frontend_origins = os.getenv("FRONTEND_ORIGINS", "*") allow_origins = ( ["*"] if frontend_origins == "*" else [origin.strip() for origin in frontend_origins.split(",") if origin.strip()] ) app.add_middleware( CORSMiddleware, allow_origins=allow_origins, allow_credentials=False, allow_methods=["*"], allow_headers=["*"], ) # Where annotations are written. Kept separate from HF_VIDEO_DATASET_REPO (the # source videos, below) so the two can never be confused with each other. HF_ANNOTATIONS_DATASET_REPO = os.getenv( "HF_ANNOTATIONS_DATASET_REPO", "dghadiya/t2av_eval_annotations", ) HF_TOKEN = os.getenv("HF_TOKEN") HF_DATASET_PRIVATE = os.getenv("HF_ANNOTATIONS_DATASET_PRIVATE", "true").lower() in { "1", "true", "yes", } ANNOTATIONS_DIR = "annotations" # Balanced-assignment records live in their own prefixes in the same repo - # annotations/ is never written to by any assignment/completion code path. # Layout: assignments/{annotator_id}/current.json (pointer to the active # round), assignments/{annotator_id}/rounds/{assignment_id}.json (immutable # per-round record), completions/{annotator_id}/{assignment_id}.json # (immutable, written once a round is verified fully complete). ASSIGNMENTS_DIR = "assignments" COMPLETIONS_DIR = "completions" api = HfApi() # Where the evaluation videos live (public dataset, read-only). HF_VIDEO_DATASET_REPO = os.getenv("HF_VIDEO_DATASET_REPO", "dghadiya/T2AV_TemporalConstraint") VIDEO_MODELS = ["LTX_2p3", "MOVA"] # Read from backend/data (not frontend/src) - the Dockerfile only copies # frontend/dist (the built output) into the final image, not frontend/src, so # a path under frontend/src would not exist at runtime in the deployed # container. backend/ is copied whole, so this path always exists there. # frontend/src/data/prompts.json is a separate copy bundled into the JS build; # keep both in sync when the prompt set changes. PROMPTS_PATH = Path(__file__).resolve().parent / "data" / "prompts.json" _video_catalog_cache = None # Guards the read-existing -> compute-balance -> write critical section in # get_or_create_assignment() so two annotators arriving at the same moment # can't both read the same low-coverage snapshot and double-assign scarce # videos. This is a single-process lock: it does not protect across multiple # Uvicorn workers/replicas, which this deployment does not use (see # Dockerfile - no --workers flag). _assignment_lock = threading.Lock() # Short-lived cache for list_annotation_records(), used only by the # assignment endpoints (lookup/create/admin summary), which can tolerate a # few seconds of staleness. The /annotations and /annotations.csv endpoints # intentionally do NOT use this cache and always return fresh data. _annotation_records_cache = {"data": None, "fetched_at": 0.0} ANNOTATION_RECORDS_CACHE_TTL_SECONDS = 15.0 class AnnotationPayload(BaseModel): video_id: str annotations: dict user: Optional[str] = "anonymous" @app.on_event("startup") def startup(): ensure_dataset_configured() validate_prompt_data() try: fetch_video_catalog() except Exception as exc: print(f"Warning: failed to prefetch video catalog at startup: {exc}") @app.get("/health") def root(): return { "status": "running", "annotation_store": "huggingface_dataset", "dataset_repo": HF_ANNOTATIONS_DATASET_REPO, } def ensure_dataset_configured(): if not HF_ANNOTATIONS_DATASET_REPO: raise HTTPException( status_code=500, detail="HF_ANNOTATIONS_DATASET_REPO is not configured.", ) def ensure_write_token(): if not HF_TOKEN: raise HTTPException( status_code=500, detail="HF_TOKEN is not configured. Add it as a Hugging Face Space secret.", ) def ensure_dataset_exists(): ensure_dataset_configured() ensure_write_token() api.create_repo( repo_id=HF_ANNOTATIONS_DATASET_REPO, repo_type="dataset", private=HF_DATASET_PRIVATE, token=HF_TOKEN, exist_ok=True, ) def slugify(value): slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value).strip("-") return slug or "unknown" def annotation_path(annotation): created_at = datetime.fromisoformat(annotation["created_at"].replace("Z", "+00:00")) timestamp = created_at.strftime("%Y%m%dT%H%M%S%fZ") video_id = slugify(annotation["video_id"]) user = slugify(annotation["user"]) return f"{ANNOTATIONS_DIR}/{timestamp}_{video_id}_{user}_{annotation['id']}.json" def read_annotation_file(path): try: local_path = hf_hub_download( repo_id=HF_ANNOTATIONS_DATASET_REPO, filename=path, repo_type="dataset", token=HF_TOKEN, ) except Exception as exc: raise HTTPException(status_code=502, detail="Failed to read annotation file") from exc with open(local_path, encoding="utf-8") as annotation_file: return json.load(annotation_file) def list_annotation_paths(): ensure_dataset_configured() try: files = api.list_repo_files( repo_id=HF_ANNOTATIONS_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, ) except RepositoryNotFoundError: # Fresh deployment: the dataset repo doesn't exist until the first # save_annotation call creates it - no repo means no annotations # exist yet, not an error (also relied on by the assignment feature, # which needs this to succeed before anyone has ever saved anything). return [] except Exception as exc: raise HTTPException(status_code=502, detail="Failed to list annotation dataset") from exc return sorted( [ path for path in files if path.startswith(f"{ANNOTATIONS_DIR}/") and path.endswith(".json") ], reverse=True, ) def list_annotation_records(): annotations = [read_annotation_file(path) for path in list_annotation_paths()] return sorted( annotations, key=lambda item: item.get("created_at", ""), reverse=True, ) def list_annotation_records_cached(): """Same data as list_annotation_records(), refreshed at most every ANNOTATION_RECORDS_CACHE_TTL_SECONDS. Used only by the assignment endpoints, which can tolerate brief staleness in exchange for not re-downloading every annotation file on every debounced keystroke.""" now = time.monotonic() stale = ( _annotation_records_cache["data"] is None or (now - _annotation_records_cache["fetched_at"]) > ANNOTATION_RECORDS_CACHE_TTL_SECONDS ) if stale: _annotation_records_cache["data"] = list_annotation_records() _annotation_records_cache["fetched_at"] = now return _annotation_records_cache["data"] UPLOAD_MAX_ATTEMPTS = 3 UPLOAD_BACKOFF_SECONDS = [0.5, 1, 2] def _is_transient_upload_error(exc): status_code = getattr(getattr(exc, "response", None), "status_code", None) if status_code is None: # Connection errors, timeouts, etc. have no HTTP status - treat as transient. return True return status_code == 429 or 500 <= status_code < 600 def upload_with_retry(**kwargs): """Uploads a file to the HF dataset repo (annotations or assignments), retrying transient failures (network errors, 429, 5xx) with exponential backoff. Non-transient errors (e.g. an auth/permission problem) are raised immediately instead of being retried.""" last_exc = None for attempt in range(UPLOAD_MAX_ATTEMPTS): try: return api.upload_file(**kwargs) except Exception as exc: last_exc = exc is_last_attempt = attempt == UPLOAD_MAX_ATTEMPTS - 1 if is_last_attempt or not _is_transient_upload_error(exc): raise time.sleep(UPLOAD_BACKOFF_SECONDS[attempt]) raise last_exc @app.post("/save_annotation") def save_annotation(payload: AnnotationPayload): ensure_dataset_exists() created_at = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") annotation = { "id": uuid4().hex, "video_id": payload.video_id, "user": payload.user or "anonymous", "annotations": payload.annotations, "created_at": created_at, } path_in_repo = annotation_path(annotation) annotation_bytes = json.dumps(annotation, indent=2).encode("utf-8") try: upload_with_retry( path_or_fileobj=io.BytesIO(annotation_bytes), path_in_repo=path_in_repo, repo_id=HF_ANNOTATIONS_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, commit_message=f"Add annotation {annotation['id']}", ) except Exception as exc: raise HTTPException(status_code=500, detail="Failed to save annotation after retries") from exc return {"message": "saved", "annotation": annotation, "path": path_in_repo} @app.get("/annotations") def list_annotations(): return list_annotation_records() @app.get("/annotations.csv") def download_annotations_csv(): annotations = list_annotation_records() output = io.StringIO() writer = csv.DictWriter( output, fieldnames=["id", "video_id", "user", "created_at", "annotations"], ) writer.writeheader() for annotation in annotations: writer.writerow( { "id": annotation["id"], "video_id": annotation["video_id"], "user": annotation["user"], "created_at": annotation["created_at"], "annotations": json.dumps(annotation["annotations"]), } ) output.seek(0) return StreamingResponse( iter([output.getvalue()]), media_type="text/csv", headers={"Content-Disposition": "attachment; filename=annotations.csv"}, ) @app.get("/annotations/{annotation_id}") def get_annotation(annotation_id: str): for annotation in list_annotation_records(): if annotation["id"] == annotation_id: return annotation raise HTTPException(status_code=404, detail="Annotation not found") VARIANT_PATTERN = re.compile(r"^main_videos/([^/]+)/([^/]+)/.+_v(\d+)_raw\.mp4$") def load_prompt_items(): with open(PROMPTS_PATH, encoding="utf-8") as prompts_file: return json.load(prompts_file) def validate_prompt_data(): """Fails fast at startup if prompts.json violates the event-id schema check_items() enforces (see validate_prompts.py). A violation doesn't raise anywhere else - it silently corrupts the annotation screen (merged annotation state, garbled temporal-constraint text) instead - so this is the only place that catches it before annotators do.""" errors = check_items(load_prompt_items()) if errors: raise RuntimeError( "prompts.json failed schema validation:\n" + "\n".join(f" - {e}" for e in errors) ) def fetch_video_catalog(): """Lists every video under main_videos/ once and caches it in memory as {model: {item_id: [(variant_number, path), ...sorted by variant]}}. Raises on failure so the startup hook can log it and callers can retry.""" global _video_catalog_cache files = api.list_repo_files( repo_id=HF_VIDEO_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, ) catalog = {model: {} for model in VIDEO_MODELS} for path in files: match = VARIANT_PATTERN.match(path) if not match: continue model, item_id, variant = match.groups() if model not in catalog: continue catalog[model].setdefault(item_id, []).append((int(variant), path)) for model in catalog: for item_id in catalog[model]: catalog[model][item_id].sort(key=lambda pair: pair[0]) _video_catalog_cache = catalog return catalog def get_video_catalog(): if _video_catalog_cache is None: return fetch_video_catalog() return _video_catalog_cache @app.get("/videos") def list_videos(): try: catalog = get_video_catalog() except Exception as exc: raise HTTPException( status_code=502, detail=f"Failed to load video catalog from {HF_VIDEO_DATASET_REPO}: {exc}", ) from exc items = [] missing = [] for prompt in load_prompt_items(): item_id = prompt["item_id"] for model in VIDEO_MODELS: variants = catalog.get(model, {}).get(item_id, []) if not variants: missing.append( { "itemId": item_id, "model": model, "reason": f"No .mp4 found under main_videos/{model}/{item_id}/ in {HF_VIDEO_DATASET_REPO}", } ) continue variant_number, path = variants[0] items.append( { "itemId": item_id, "model": model, "variant": f"v{variant_number}", "videoUrl": hf_hub_url( repo_id=HF_VIDEO_DATASET_REPO, repo_type="dataset", filename=path, ), "path": path, } ) return {"items": items, "missing": missing} def _current_catalog_video_ids(): """The same (itemId, model) pairs list_videos() resolves to real files, reduced to just the video_id strings used everywhere else (save_annotation, assignments) - i.e. only videos that are actually playable right now are eligible for assignment.""" catalog = get_video_catalog() ids = [] for prompt in load_prompt_items(): item_id = prompt["item_id"] for model in VIDEO_MODELS: if catalog.get(model, {}).get(item_id): ids.append(f"{item_id}__{model}") return ids def _round_pointer_path(annotator_id): return f"{ASSIGNMENTS_DIR}/{annotator_id}/current.json" def _round_path(annotator_id, assignment_id): return f"{ASSIGNMENTS_DIR}/{annotator_id}/rounds/{assignment_id}.json" def _completion_path(annotator_id, assignment_id): return f"{COMPLETIONS_DIR}/{annotator_id}/{assignment_id}.json" def _read_json_from_repo(path): """Returns the parsed JSON at `path` in the annotations dataset repo, or None if it doesn't exist (distinct from a real read failure, which still raises). Missing repo (e.g. a fresh deployment before the first save_annotation/get_or_create_current_round call has ever created it) is treated the same as a missing file - both mean "there's nothing here yet" to every caller.""" try: local_path = hf_hub_download( repo_id=HF_ANNOTATIONS_DATASET_REPO, filename=path, repo_type="dataset", token=HF_TOKEN, ) except (EntryNotFoundError, RepositoryNotFoundError): return None except Exception as exc: raise HTTPException(status_code=502, detail=f"Failed to read {path}") from exc with open(local_path, encoding="utf-8") as json_file: return json.load(json_file) def _write_json_to_repo(path, data, commit_message): data_bytes = json.dumps(data, indent=2).encode("utf-8") try: upload_with_retry( path_or_fileobj=io.BytesIO(data_bytes), path_in_repo=path, repo_id=HF_ANNOTATIONS_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, commit_message=commit_message, ) except Exception as exc: raise HTTPException(status_code=500, detail=f"Failed to save {path} after retries") from exc def list_active_round_records(): """The CURRENT round for every annotator who has one, excluding rounds that already have a completion record - a finished round's videos are already fully reflected in completed_by_video, so they're not an open "reservation" against other annotators' balancing anymore. Used only for reserved-by-video coverage, not for anything correctness-critical (a returning annotator's own round is always read directly by pointer, not found via this list).""" ensure_dataset_configured() try: files = api.list_repo_files( repo_id=HF_ANNOTATIONS_DATASET_REPO, repo_type="dataset", token=HF_TOKEN, ) except RepositoryNotFoundError: return [] except Exception as exc: raise HTTPException(status_code=502, detail="Failed to list assignment records") from exc pointer_paths = sorted( path for path in files if path.startswith(f"{ASSIGNMENTS_DIR}/") and path.endswith("/current.json") ) records = [] for pointer_path in pointer_paths: annotator_id = pointer_path.split("/")[1] pointer = _read_json_from_repo(pointer_path) if pointer is None: continue round_record = _read_json_from_repo(_round_path(annotator_id, pointer["assignment_id"])) if round_record is None: continue if _read_json_from_repo(_completion_path(annotator_id, pointer["assignment_id"])) is not None: continue records.append(round_record) return records def _completed_maps(fresh=False): """`fresh=True` bypasses the 15s annotation-records cache. Used only for the completion check in get_or_create_current_round(), where staleness would visibly delay an annotator's own completion code right after their 20th save - every other caller (lookup, round-balancing, admin summary) can tolerate a few seconds of lag.""" annotation_records = list_annotation_records() if fresh else list_annotation_records_cached() deduped, skipped = assignment.dedupe_latest(annotation_records) return ( deduped, skipped, assignment.completed_by_video(deduped), assignment.completed_by_annotator(deduped), ) def _round_response(round_record, completion, saved_responses=None): response = { "annotator_id": round_record["annotator_id"], "assignment_id": round_record["assignment_id"], "round_number": round_record["round_number"], "video_ids": round_record["video_ids"], "requested_size": round_record["requested_size"], "actual_size": round_record["actual_size"], "status": "completed" if completion else "in_progress", } if completion: response["completion_code"] = completion["completion_code"] response["completed_at"] = completion["completed_at"] else: # Lets a returning annotator (same name, new session or a reload) # resume with prior answers pre-filled instead of the frontend # starting every field blank - see EvaluationPage.jsx's initial # state, which seeds from this. response["saved_annotations"] = saved_responses or {} return response def _persist_completion(annotator_id, round_record): completion = { "annotator_id": annotator_id, "assignment_id": round_record["assignment_id"], "round_number": round_record["round_number"], "completion_code": assignment.generate_completion_code(), "completed_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"), } ensure_dataset_exists() _write_json_to_repo( _completion_path(annotator_id, round_record["assignment_id"]), completion, f"Complete round {round_record['round_number']} for {annotator_id}", ) return completion def _create_round(annotator_id, annotator_raw, round_number): catalog_video_ids = _current_catalog_video_ids() _deduped, _skipped, completed_by_video_map, completed_by_annotator_map = _completed_maps() already_completed = completed_by_annotator_map.get(annotator_id, set()) active_rounds = list_active_round_records() reserved_map = assignment.reserved_by_video( active_rounds, completed_by_annotator_map, datetime.now(timezone.utc), assignment.RESERVATION_TTL_SECONDS, exclude_annotator=annotator_id, ) round_record = assignment.build_round( annotator_id, annotator_raw, catalog_video_ids, already_completed, completed_by_video_map, reserved_map, random.Random(), round_number, ) round_record["assignment_id"] = uuid4().hex round_record["created_at"] = datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") ensure_dataset_exists() _write_json_to_repo( _round_path(annotator_id, round_record["assignment_id"]), round_record, f"Create round {round_number} for {annotator_id}", ) _write_json_to_repo( _round_pointer_path(annotator_id), {"assignment_id": round_record["assignment_id"], "round_number": round_number}, f"Point {annotator_id} at round {round_number}", ) # A fresh round's video_ids are, by construction, videos this annotator # has never completed (build_round excludes already_completed from the # candidate pool) - nothing to rehydrate yet, but saved_responses is # always present in the response shape either way. return _round_response(round_record, completion=None, saved_responses={}) def get_or_create_current_round(annotator_raw): """The full POST /assignment state machine - one lock-guarded decision per call, no separate "start a new round" action anywhere: - no current round yet -> create round 1 - current round not completion-verified -> re-check against real annotation data; still incomplete -> return it unchanged (same 20); just became fully complete -> generate + persist the completion record now, then fall into the completed branch below - current round completion-verified, within the grace period -> return the same completed round + code every time (refresh-safe) - current round completion-verified, past the grace period -> create and return the next round automatically - this is the only "advance to a new round" trigger, no button/endpoint needed """ ensure_dataset_configured() annotator_id = assignment.slugify(annotator_raw) with _assignment_lock: pointer = _read_json_from_repo(_round_pointer_path(annotator_id)) if pointer is None: return _create_round(annotator_id, annotator_raw, round_number=1) round_record = _read_json_from_repo(_round_path(annotator_id, pointer["assignment_id"])) if round_record is None: raise HTTPException( status_code=500, detail="Assignment pointer is inconsistent with stored round data", ) completion = _read_json_from_repo(_completion_path(annotator_id, pointer["assignment_id"])) if completion is None: deduped, _skipped, _completed_by_video_map, completed_by_annotator_map = _completed_maps(fresh=True) completed_set = completed_by_annotator_map.get(annotator_id, set()) all_done = len(round_record["video_ids"]) > 0 and all( video_id in completed_set for video_id in round_record["video_ids"] ) if not all_done: saved_responses = assignment.saved_responses_for_round( deduped, annotator_id, round_record["video_ids"] ) return _round_response(round_record, completion=None, saved_responses=saved_responses) completion = _persist_completion(annotator_id, round_record) completed_at = assignment.parse_iso(completion["completed_at"]) seconds_since_completion = (datetime.now(timezone.utc) - completed_at).total_seconds() if seconds_since_completion <= assignment.COMPLETION_GRACE_PERIOD_SECONDS: return _round_response(round_record, completion=completion) return _create_round(annotator_id, annotator_raw, round_number=round_record["round_number"] + 1) def lookup_assignment(annotator_raw): """Read-only, advisory check for the frontend's name-collision guard. Never creates or advances anything - reports the CURRENT round's progress if one exists. Response shape is unchanged from before.""" ensure_dataset_configured() annotator_id = assignment.slugify(annotator_raw) pointer = _read_json_from_repo(_round_pointer_path(annotator_id)) if pointer is None: return {"exists": False} round_record = _read_json_from_repo(_round_path(annotator_id, pointer["assignment_id"])) if round_record is None: return {"exists": False} video_ids = round_record.get("video_ids", []) _deduped, _skipped, _completed_by_video_map, completed_by_annotator_map = _completed_maps() completed = completed_by_annotator_map.get(annotator_id, set()) return { "exists": True, "completed": len(completed & set(video_ids)), "total": len(video_ids), } class AssignmentRequest(BaseModel): annotator: str @app.post("/assignment") def create_or_get_assignment(payload: AssignmentRequest): annotator_raw = (payload.annotator or "").strip() if not annotator_raw: raise HTTPException(status_code=400, detail="annotator is required") return get_or_create_current_round(annotator_raw) @app.get("/assignment/lookup") def assignment_lookup(annotator: str = ""): annotator_raw = annotator.strip() if not annotator_raw: return {"exists": False} return lookup_assignment(annotator_raw) @app.get("/admin/summary") def admin_summary(): catalog_video_ids = _current_catalog_video_ids() annotation_records = list_annotation_records_cached() deduped, skipped, completed_by_video_map, completed_by_annotator_map = _completed_maps() active_rounds = list_active_round_records() reserved_map = assignment.reserved_by_video( active_rounds, completed_by_annotator_map, datetime.now(timezone.utc), assignment.RESERVATION_TTL_SECONDS, ) duplicate_records_deduped = len(annotation_records) - skipped - len(deduped) return assignment.summarize( catalog_video_ids, completed_by_video_map, reserved_map, total_completed_preserved=len(deduped), skipped_records=skipped, duplicate_records_deduped=duplicate_records_deduped, assignment_count=len(active_rounds), ) static_dir = Path(__file__).resolve().parent.parent / "frontend" / "dist" if static_dir.exists(): app.mount("/", StaticFiles(directory=static_dir, html=True), name="frontend")