"""Fixed-group assignment and legacy balanced assignment: pure logic, no FastAPI/HF imports. Everything here takes plain data in and returns plain data out, so it can be unit-tested 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. Legacy assignment settings are env-overridable, matching the existing HF_* variable convention in main.py. Fixed-group size and slots per group are function parameters, defaulting to at most 50 videos and 3 participants. Saved groups are reused rather than rebuilt when these defaults change. """ import os import re import secrets import string from datetime import datetime # Legacy round configuration; new groups default to at most 50 videos and 3 slots. VIDEOS_PER_ANNOTATOR = int(os.getenv("VIDEOS_PER_ANNOTATOR", "1000")) TARGET_ANNOTATIONS_PER_VIDEO = int( os.getenv("TARGET_ANNOTATIONS_PER_VIDEO", "3") ) # Abandonment TTL for both legacy reservations and fixed-group assignments # that have not saved any answers. RESERVATION_TTL_SECONDS = int(os.getenv("ASSIGNMENT_RESERVATION_TTL_SECONDS", str(24 * 3600))) COMPLETION_CODE_ALPHABET = string.ascii_uppercase + string.digits COMPLETION_CODE_LENGTH = 12 def build_allocation_state(catalog_video_ids, group_size=50, slots_per_group=3): """Build initial allocation_state.json data: groups and their member slots. Uses actual resolved video IDs rather than a configured catalog size. Call only when no saved allocation state exists; reuse it thereafter. Members map annotator IDs to in_progress/completed status. Existing assignments continue using the legacy functions below. """ for name, value in (("group_size", group_size), ("slots_per_group", slots_per_group)): if type(value) is not int or value <= 0: raise ValueError(f"{name} must be a positive integer") ids = list(catalog_video_ids) if not ids or any(not isinstance(v, str) or not v.strip() for v in ids): raise ValueError("catalog must contain non-empty video IDs") if len(ids) != len(set(ids)): raise ValueError("catalog contains duplicate video IDs") group_count = (len(ids) + group_size - 1) // group_size smaller_size, larger_group_count = divmod(len(ids), group_count) group_sizes = [smaller_size + 1] * larger_group_count group_sizes += [smaller_size] * (group_count - larger_group_count) groups = [] start = 0 for index, size in enumerate(group_sizes, start=1): groups.append({ "group_id": f"group_{index:04d}", "video_ids": ids[start:start + size], "members": {}, }) start += size return { "schema_version": 1, "group_size": group_size, "slots_per_group": slots_per_group, "total_videos": len(ids), "groups": groups, } def build_workload_allocation_state(group_specs, slots_per_group=2): """Build persisted variable-size groups for the workload-duration pilot.""" if not isinstance(group_specs, list) or not group_specs: raise ValueError("group_specs must be a non-empty list") if type(slots_per_group) is not int or slots_per_group <= 0: raise ValueError("slots_per_group must be a positive integer") groups = [] all_ids = [] group_ids = set() for spec in group_specs: if not isinstance(spec, dict): raise ValueError("each group spec must be an object") group_id = spec.get("group_id") video_ids = spec.get("video_ids") if (not isinstance(group_id, str) or not group_id or not isinstance(video_ids, list) or not video_ids or any(not isinstance(v, str) or not v.strip() for v in video_ids)): raise ValueError("each group must have an ID and non-empty video IDs") if group_id in group_ids: raise ValueError("workload group IDs must be unique") group_ids.add(group_id) if len(video_ids) != len(set(video_ids)): raise ValueError(f"group {group_id} contains duplicate video IDs") all_ids.extend(video_ids) groups.append({ "group_id": group_id, "video_ids": list(video_ids), "members": {}, }) if len(all_ids) != len(set(all_ids)): raise ValueError("a video cannot appear in more than one workload group") return { "schema_version": 1, "allocation_mode": "workload_test", "group_size": None, "group_sizes": [len(group["video_ids"]) for group in groups], "slots_per_group": slots_per_group, "total_videos": len(all_ids), "groups": groups, } def pick_available_group(state, rng): """Choose a group ID from saved state, or None if every group is full. Both in-progress and completed members occupy slots. Does not mutate state or claim a slot; the I/O caller must persist claims atomically and re-read/recompute after a commit conflict. """ available = [ group for group in state["groups"] if len(group["members"]) + group.get("legacy_reserved_slots", 0) < state["slots_per_group"] ] if not available: return None lowest = min(len(group["members"]) + group.get("legacy_reserved_slots", 0) for group in available) return rng.choice([ group["group_id"] for group in available if len(group["members"]) + group.get("legacy_reserved_slots", 0) == lowest ]) 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 build_user_state(round_record, latest_annotations=None, completion=None): """Build users/{annotator_id}.json from an assignment and latest records. latest_annotations maps video IDs to metadata (id, created_at, annotation_path), never full answer bodies. This also supports migration of existing rounds; group_id is None for legacy assignments. """ import copy latest = copy.deepcopy(latest_annotations or {}) video_ids = set(round_record["video_ids"]) if not set(latest) <= video_ids: raise ValueError("annotation video is outside the current assignment") for record in latest.values(): if not all(record.get(field) for field in ("id", "created_at", "annotation_path")): raise ValueError("latest annotation metadata is incomplete") if completion is not None: if (completion["annotator_id"] != round_record["annotator_id"] or completion["assignment_id"] != round_record["assignment_id"]): raise ValueError("completion belongs to a different assignment") if not video_ids or set(latest) != video_ids: raise ValueError("cannot complete a round with unfinished videos") return { "schema_version": 1, "annotator_id": round_record["annotator_id"], "group_id": round_record.get("group_id"), "assignment": copy.deepcopy(round_record), "completed_video_ids": sorted(latest), "latest_annotations": latest, "status": "completed" if completion is not None else "in_progress", "completion_code": completion["completion_code"] if completion is not None else None, "completed_at": completion["completed_at"] if completion is not None else None, } def normalize_participant_id(value): """Trim input without silently replacing invalid characters or merging IDs. Keep legacy canonical names usable; this is syntax validation, not verification of a Prolific account. """ value = value.strip() if isinstance(value, str) else "" if not re.fullmatch(r"[A-Za-z0-9_][A-Za-z0-9_.-]{0,127}", value): raise ValueError("Enter your Prolific ID without spaces or special characters") return value.lower() 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 fixed_assignment_is_abandoned(state, now, ttl_seconds=RESERVATION_TTL_SECONDS): """Return whether an untouched fixed-group assignment may release its slot. Any saved answer preserves the assignment indefinitely. Completed and legacy assignments are never expired by this rule. """ if (not state or state.get("group_id") is None or state.get("status") != "in_progress" or state.get("completed_video_ids") or state.get("latest_annotations")): return False try: created_at = parse_iso(state["assignment"]["created_at"]) except (KeyError, ValueError, AttributeError, TypeError): return False return (now - created_at).total_seconds() > ttl_seconds 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 up to `count` videos, lowest coverage (completed + reserved) first, random tie-break among equal coverage. A video already at `target` is never picked. Returns fewer than `count` when fewer under-target videos remain.""" if count <= 0: return [] under_target = [v for v in candidates if coverage_fn(v) < target] ordered = sorted(under_target, 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: up to `videos_per_annotator` videos this annotator has never completed before. Videos already completed are excluded. Videos already at the coverage target are excluded too, so a short round is returned instead of giving a video a fourth rater.""" 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, }