submit.py

"""Submission handling for the validation phase.

Validates a participant's predictions.jsonl + metadata, enforces the per-user
daily rate limit and the division param cap, then uploads request.json +
predictions.jsonl to the requests dataset with status=PENDING. The eval worker
takes it from there.

The Space identifies the participant via HF OAuth (hf_user) but uploads with the
Space's own write token, because the requests dataset is org-owned and a
participant's OAuth token has no write access to it.
"""

from __future__ import annotations

import io
import json
import random
import string
from datetime import datetime, timezone

from huggingface_hub import HfApi, hf_hub_download

import config


def _now_iso() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _today() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%d")


def _new_sid() -> str:
    stamp = datetime.now(timezone.utc).strftime("%Y-%m-%dT%H-%M-%S")
    suffix = "".join(random.choices(string.ascii_lowercase + string.digits, k=6))
    return f"{stamp}_{suffix}"


def parse_predictions(raw: str, track: str) -> tuple[list[dict], dict, str]:
    """Parse + validate the JSONL body. Returns (rows, self_scores, error).

    The file is the single source of self-reported scores: every line is a JSON
    object that is either a prediction row (has "video_path") or a single
    metadata line carrying the self-reported score(s) (has "llm_judge" and no
    "video_path", e.g. `{"llm_judge": 0.83}` for EgoConv). error=='' on ok.
    """
    rows: list[dict] = []
    self_scores: dict = {}
    for i, line in enumerate(raw.splitlines(), 1):
        line = line.strip()
        if not line:
            continue
        try:
            obj = json.loads(line)
        except json.JSONDecodeError as e:
            return [], {}, f"line {i}: invalid JSON ({e})"
        if not isinstance(obj, dict):
            return [], {}, f"line {i}: expected a JSON object"
        if "video_path" in obj:
            rows.append(obj)
        elif "llm_judge" in obj:
            if self_scores:
                return [], {}, (f"line {i}: more than one self-reported-score line "
                                "(include exactly one)")
            self_scores = {"llm_judge": obj["llm_judge"]}
        else:
            return [], {}, (f"line {i}: object must contain 'video_path' (a prediction) "
                            "or 'llm_judge' (the self-reported score)")

    if len(rows) != config.N_GOLD:
        return [], {}, f"expected exactly {config.N_GOLD} prediction rows, got {len(rows)}"

    required = config.PREDICTION_KEYS[track]
    seen: set[str] = set()
    for i, obj in enumerate(rows, 1):
        missing = required - set(obj.keys())
        if missing:
            return [], {}, f"row {i}: missing keys {sorted(missing)}"
        vp = obj.get("video_path")
        if vp in seen:
            return [], {}, f"row {i}: duplicate video_path {vp!r}"
        seen.add(vp)

    # Coverage: the prediction ids must match the gold set exactly (if bundled).
    gold = config.load_gold_ids(track)
    if gold is not None and seen != gold:
        extra = sorted(seen - gold)[:3]
        missing_ids = sorted(gold - seen)[:3]
        return [], {}, (
            f"video_path set does not match the {len(gold)} gold ids "
            f"(unexpected e.g. {extra}; missing e.g. {missing_ids})"
        )
    return rows, self_scores, ""


def count_user_today(api: HfApi, hf_user: str, track: str, division: str) -> int:
    """Submissions by this user to this subtrack (track, division) today (UTC).

    Per the challenge rules, the val limit is per "Subtrack" = (track, division)
    per UTC day. Reads each of today's request.json for the subtrack to bind to
    the user. Volume is low in the val phase (a handful/team/day), so the extra
    reads are cheap.
    """
    today = _today()
    prefix = f"requests/{config.PHASE}/{track}/{division}/"
    try:
        files = api.list_repo_files(config.REQUESTS_REPO, repo_type="dataset")
    except Exception:
        return 0

    count = 0
    for f in files:
        if not (f.startswith(prefix) and f.endswith("/request.json")):
            continue
        sid = f.split("/")[-2]  # .../<division>/<sid>/request.json
        if sid[:10] != today:  # sid = YYYY-MM-DDThh-mm-ss_xxxxxx
            continue
        try:
            local = hf_hub_download(
                config.REQUESTS_REPO, f, repo_type="dataset", token=api.token
            )
            with open(local) as fh:
                if json.load(fh).get("hf_user") == hf_user:
                    count += 1
        except Exception:
            continue  # unreadable request shouldn't block a legit submit
    return count


def validate_only(
    *,
    track: str,
    division: str,
    total_params: int,
    active_params: int,
    predictions_raw: str,
) -> tuple[bool, str]:
    """All format / eligibility checks, EXCEPT the daily rate-limit and upload.

    Shared by the real submit path and the Submit tab's "Validate (no submit)"
    button, so the dry-run check is byte-identical to what a real submit enforces.
    The self-reported score (EgoConv) is read from the predictions file itself.
    Returns (ok, message_for_user).
    """
    if track not in config.TRACKS:
        return False, f"Unknown track: {track}"
    if division not in config.DIVISIONS:
        return False, f"Unknown division: {division}"

    # Division param cap (self-declared in val).
    declared_div = config.division_for_params(int(total_params))
    if declared_div is None:
        return False, f"total_params={total_params} must be a positive parameter count."
    if declared_div != division:
        return False, (
            f"total_params={total_params / 1e9:.2f}B falls in '{declared_div}' "
            f"but you selected '{division}'."
        )
    if not (0 < int(active_params) <= int(total_params)):
        return False, "active_params must be > 0 and <= total_params."

    # Predictions format + gold coverage; the self-reported score comes from the file.
    rows, self_scores, why = parse_predictions(predictions_raw, track)
    if why:
        return False, f"predictions.jsonl invalid: {why}"

    # Self-reported scores (ConvQA: LLM-Judge required; other tracks: none).
    ok, why = config.validate_self_report(track, self_scores)
    if not ok:
        if track == "convqa":
            return False, (
                f"{why}. EgoConv ranks on a self-reported LLM-Judge score: add a line "
                f'`{{"llm_judge": <score 0-1>}}` to your predictions.jsonl, computed with '
                f"{config.CONVQA_JUDGE_MODEL} via the starter kit (see the About tab)."
            )
        return False, why

    extra = ""
    if track == "convqa":
        extra = f", self-reported LLM-Judge={float(self_scores['llm_judge']):.3f}"
    return True, (
        f"Format valid: {len(rows)} rows, ids match the gold set, keys OK for "
        f"{config.TRACK_LABELS[track]} / {division}{extra}."
    )


def validate_and_submit(
    *,
    hf_user: str,
    track: str,
    division: str,
    team_name: str,
    model_name: str,
    license_str: str,
    open_weight: bool,
    total_params: int,
    active_params: int,
    predictions_raw: str,
    token: str,
) -> tuple[bool, str]:
    """Full submit path. Returns (ok, message_for_user)."""
    # First, and before anything is read or written. A closed phase is not a
    # validation failure, so it is not reported through validate_only(): a team
    # whose file is fine should not be told to fix it.
    if not config.VAL_PHASE_OPEN:
        return False, config.VAL_CLOSED_NOTICE
    if not team_name.strip():
        return False, "Team name is required."
    if not token:
        return False, "Server is missing its write token; contact the organizers."

    ok, why = validate_only(
        track=track,
        division=division,
        total_params=int(total_params),
        active_params=int(active_params),
        predictions_raw=predictions_raw,
    )
    if not ok:
        return False, why

    # Self-reported score is carried IN the predictions file (already validated above).
    _, self_reported_scores, _ = parse_predictions(predictions_raw, track)

    api = HfApi(token=token)

    # Rate limit: per (hf_user, track, division) = "Subtrack" per UTC day (challenge rules).
    if count_user_today(api, hf_user, track, division) >= config.MAX_SUBMISSIONS_PER_DAY:
        return False, (
            f"Daily limit reached: {config.MAX_SUBMISSIONS_PER_DAY} submissions/"
            f"day/subtrack for {hf_user} on {config.TRACK_LABELS[track]} ({division})."
        )

    sid = _new_sid()
    base = f"requests/{config.PHASE}/{track}/{division}/{sid}"
    request = {
        "submission_id": sid,
        "phase": config.PHASE,
        "track": track,
        "division": division,
        "team_name": team_name.strip(),
        "hf_user": hf_user,
        "model_name": model_name.strip(),
        "license": license_str.strip(),
        "open_weight": bool(open_weight),
        "total_params": int(total_params),
        "active_params": int(active_params),
        # Self-reported scores: ConvQA carries a participant-computed LLM-Judge
        # score (the val-phase ranked metric, badged self-reported); the worker
        # still recomputes verified BLEU as a cross-check. Other tracks are fully
        # organizer-scored, so this is empty. Latency is collected only in the
        # test phase (organizer runs every model on the same hardware).
        "self_reported_scores": {
            k: float(v) for k, v in (self_reported_scores or {}).items()
        },
        "predictions_path": f"{base}/predictions.jsonl",
        "created_at": _now_iso(),
        "status": "PENDING",
    }

    try:
        api.upload_file(
            path_or_fileobj=io.BytesIO(json.dumps(request, indent=2).encode()),
            path_in_repo=f"{base}/request.json",
            repo_id=config.REQUESTS_REPO,
            repo_type="dataset",
            commit_message=f"submit {sid} ({track}/{division}) by {hf_user}",
        )
        api.upload_file(
            path_or_fileobj=io.BytesIO(predictions_raw.encode()),
            path_in_repo=f"{base}/predictions.jsonl",
            repo_id=config.REQUESTS_REPO,
            repo_type="dataset",
            commit_message=f"predictions for {sid}",
        )
    except Exception as e:
        # Don't surface raw exceptions (may leak internal repo paths) on a public UI.
        print(f"[submit] upload failed for {sid}: {e!r}")
        return False, "Upload failed — please retry; contact the organizers if it persists."

    return True, (
        f"✅ Submitted as `{sid}` ({config.TRACK_LABELS[track]} / {division}). "
        f"Status PENDING — the organizers' eval worker will verify and publish "
        f"your score to the leaderboard shortly."
    )