"""Balanced random subset assignment: pure logic, no FastAPI/HF imports. Everything here takes plain data in and returns plain data out, so it can be unit-tested (see test_assignment.py) without touching the network. The I/O layer (reading/writing annotation and assignment records from the HF dataset repo, locking) lives in main.py and calls into these functions. Config is env-overridable, matching the existing HF_* var convention in main.py, so these three numbers can change without touching the algorithm. """ import os import re import secrets import string from datetime import datetime TOTAL_VIDEOS = int(os.getenv("TOTAL_VIDEOS", "200")) VIDEOS_PER_ANNOTATOR = int(os.getenv("VIDEOS_PER_ANNOTATOR", "20")) TARGET_ANNOTATIONS_PER_VIDEO = int(os.getenv("TARGET_ANNOTATIONS_PER_VIDEO", "5")) # How long an in-progress (not-yet-completed) round still counts as an open # "reservation" against other annotators' balancing, before being treated as # abandoned. Distinct from COMPLETION_GRACE_PERIOD_SECONDS below. RESERVATION_TTL_SECONDS = int(os.getenv("ASSIGNMENT_RESERVATION_TTL_SECONDS", str(24 * 3600))) # How long a just-completed round keeps showing its completion code before a # later visit automatically gets the next round instead. Distinct from # RESERVATION_TTL_SECONDS above - this one is "how long do you get to read # your own code," not "when do other annotators stop waiting on your slot." COMPLETION_GRACE_PERIOD_SECONDS = int(os.getenv("COMPLETION_GRACE_PERIOD_SECONDS", str(10 * 60))) COMPLETION_CODE_ALPHABET = string.ascii_uppercase + string.digits COMPLETION_CODE_LENGTH = 12 def generate_completion_code(): """Cryptographically random completion code, format T2AV-XXXXXXXXXXXX. Uses `secrets`, not `random` - this is a credential the annotator submits elsewhere (e.g. an MTurk HIT) as proof of completion, not a UI cosmetic.""" suffix = "".join(secrets.choice(COMPLETION_CODE_ALPHABET) for _ in range(COMPLETION_CODE_LENGTH)) return f"T2AV-{suffix}" def slugify(value): """Canonical annotator_id from a raw display name. Lowercased (unlike main.py's filename slugify) so "Adi" and "adi" resolve to one identity - this is the identity key used for completion/assignment lookups, not a filename, so case-insensitivity is the more correct default here.""" slug = re.sub(r"[^A-Za-z0-9_.-]+", "-", value or "").strip("-").lower() return slug or "unknown" def parse_iso(value): return datetime.fromisoformat(value.replace("Z", "+00:00")) def dedupe_latest(annotation_records): """A "completed annotation" is one row per unique (annotator_id, video_id) pair - repeated saves/autosaves/updates from the same annotator for the same video collapse into a single record, keeping the most recent created_at. Records missing a required field are skipped (never crash on older/malformed data) and reported separately. Returns (deduped: dict[(annotator_id, video_id)] -> record, skipped: int). """ latest = {} skipped = 0 for record in annotation_records: try: video_id = record["video_id"] user_raw = record["user"] created_at = record["created_at"] if not video_id or not user_raw or not created_at: raise ValueError("empty required field") parse_iso(created_at) # validate shape before trusting string comparisons below except (KeyError, TypeError, ValueError, AttributeError): skipped += 1 continue key = (slugify(user_raw), video_id) existing = latest.get(key) if existing is None or created_at > existing["created_at"]: latest[key] = record return latest, skipped def completed_by_video(deduped): result = {} for annotator_id, video_id in deduped: result.setdefault(video_id, set()).add(annotator_id) return result def completed_by_annotator(deduped): result = {} for annotator_id, video_id in deduped: result.setdefault(annotator_id, set()).add(video_id) return result def reserved_by_video(assignment_records, completed_by_annotator_map, now, ttl_seconds, exclude_annotator=None): """Open (not-yet-completed, not-yet-expired) reservations per video, across every OTHER annotator's assignment. A reservation stops counting once its assignment is older than ttl_seconds ("abandoned") - this never deletes or mutates the stored assignment record, it only stops contributing to the live coverage snapshot used for balancing.""" result = {} for record in assignment_records: annotator_id = record.get("annotator_id") if not annotator_id or annotator_id == exclude_annotator: continue try: created_at = parse_iso(record["created_at"]) except (KeyError, ValueError, AttributeError, TypeError): continue if (now - created_at).total_seconds() > ttl_seconds: continue completed = completed_by_annotator_map.get(annotator_id, set()) for video_id in record.get("video_ids", []): if video_id in completed: continue result[video_id] = result.get(video_id, 0) + 1 return result def pick_balanced(candidates, coverage_fn, count, rng, target=TARGET_ANNOTATIONS_PER_VIDEO): """Picks `count` videos, prioritizing lowest coverage (completed + reserved) first, random tie-break among equal coverage. Prefers strictly-under-target videos; only reaches into at-or-over-target ones if there aren't enough under-target candidates to satisfy `count`.""" if count <= 0: return [] under_target = [v for v in candidates if coverage_fn(v) < target] pool = under_target if len(under_target) >= count else list(candidates) ordered = sorted(pool, key=lambda video_id: (coverage_fn(video_id), rng.random())) return ordered[:count] def build_round( annotator_id, annotator_raw, catalog_video_ids, already_completed, completed_by_video_map, reserved_by_video_map, rng, round_number, videos_per_annotator=VIDEOS_PER_ANNOTATOR, target=TARGET_ANNOTATIONS_PER_VIDEO, ): """Builds a new round: exactly `videos_per_annotator` videos this annotator has never completed before - across every prior round AND any completions that predate this feature (`already_completed` is the full, round-agnostic history). Videos already completed are unconditionally excluded, never used to seed the round. Returns fewer than `videos_per_annotator` only when the eligible pool itself is smaller (see `actual_size` vs `requested_size` in the result) - never fills the gap with already-completed videos.""" candidates = [v for v in catalog_video_ids if v not in already_completed] def coverage(video_id): return len(completed_by_video_map.get(video_id, set())) + reserved_by_video_map.get(video_id, 0) video_ids = pick_balanced(candidates, coverage, videos_per_annotator, rng, target) return { "annotator_id": annotator_id, "annotator_name_raw": annotator_raw, "round_number": round_number, "video_ids": video_ids, "requested_size": videos_per_annotator, "actual_size": len(video_ids), } def summarize( catalog_video_ids, completed_by_video_map, reserved_by_video_map, total_completed_preserved, skipped_records, duplicate_records_deduped, assignment_count, target=TARGET_ANNOTATIONS_PER_VIDEO, ): below = sum(1 for v in catalog_video_ids if len(completed_by_video_map.get(v, set())) < target) return { "total_videos": len(catalog_video_ids), "target_per_video": target, "videos_per_annotator": VIDEOS_PER_ANNOTATOR, "completed_counts": {v: len(completed_by_video_map.get(v, set())) for v in catalog_video_ids}, "reserved_counts": {v: reserved_by_video_map.get(v, 0) for v in catalog_video_ids}, "videos_below_target": below, "videos_at_or_above_target": len(catalog_video_ids) - below, "total_existing_annotations_preserved": total_completed_preserved, "skipped_malformed_records": skipped_records, "duplicate_records_deduplicated": duplicate_records_deduped, "total_assignments": assignment_count, }