src/leaderboard/sync.py

import logging
import os

from huggingface_hub import HfApi, hf_hub_download

from src.envs import (
    ON_SPACE,
    RESULTS_DATASET_REPO,
    RESULTS_DIR,
    SPACE_REPO,
    SUBMISSIONS_DIR,
    SUBMISSIONS_REPO_PREFIX,
    TOKEN,
    USE_HUB_STORAGE,
)

logger = logging.getLogger(__name__)


def _api() -> HfApi:
    return HfApi(token=TOKEN)


def log_storage_status() -> None:
    if not ON_SPACE:
        return
    if USE_HUB_STORAGE:
        print(f"Results dataset storage enabled for {RESULTS_DATASET_REPO}")
        return
    if not RESULTS_DATASET_REPO:
        print("Results storage disabled: could not resolve results dataset repo id")
        return
    print(
        f"Results storage disabled: set HF_TOKEN secret with write access to {RESULTS_DATASET_REPO}"
    )


def _ensure_results_dataset() -> None:
    if not USE_HUB_STORAGE:
        return
    try:
        _api().create_repo(
            repo_id=RESULTS_DATASET_REPO,
            repo_type="dataset",
            exist_ok=True,
            token=TOKEN,
        )
    except Exception as exc:
        logger.warning("Could not ensure results dataset exists: %s", exc)


def _download_dataset_file(path_in_repo: str) -> None:
    hf_hub_download(
        repo_id=RESULTS_DATASET_REPO,
        repo_type="dataset",
        filename=path_in_repo,
        local_dir=RESULTS_DIR,
        local_dir_use_symlinks=False,
        token=TOKEN,
    )


def _sync_legacy_space_results() -> None:
    """One-time compatibility: pull results/*.json committed to the Space repo."""
    if not SPACE_REPO or not TOKEN:
        return

    try:
        repo_files = _api().list_repo_files(repo_id=SPACE_REPO, repo_type="space")
        legacy_files = [
            path for path in repo_files
            if path.startswith("results/") and path.endswith(".json")
        ]
        for path_in_repo in legacy_files:
            filename = os.path.basename(path_in_repo)
            local_path = os.path.join(RESULTS_DIR, filename)
            if os.path.isfile(local_path):
                continue
            hf_hub_download(
                repo_id=SPACE_REPO,
                repo_type="space",
                filename=path_in_repo,
                local_dir=os.path.dirname(RESULTS_DIR),
                local_dir_use_symlinks=False,
                token=TOKEN,
            )
    except Exception as exc:
        logger.warning("Could not sync legacy Space results: %s", exc)


def sync_results_from_hub() -> None:
    """Pull submission JSON and ZIP files from the results dataset (and legacy Space repo)."""
    if not USE_HUB_STORAGE:
        return

    os.makedirs(RESULTS_DIR, exist_ok=True)
    os.makedirs(SUBMISSIONS_DIR, exist_ok=True)
    try:
        repo_files = _api().list_repo_files(repo_id=RESULTS_DATASET_REPO, repo_type="dataset")
        result_files = [
            path
            for path in repo_files
            if path.endswith(".json")
            or (
                path.startswith(f"{SUBMISSIONS_REPO_PREFIX}/")
                and path.endswith(".zip")
            )
        ]
        for path_in_repo in result_files:
            _download_dataset_file(path_in_repo)
    except Exception as exc:
        logger.exception("Could not sync results from dataset: %s", exc)

    _sync_legacy_space_results()


def upload_result_to_hub(local_path: str, model: str) -> str | None:
    """Persist a results JSON file to the results dataset (no Space restart)."""
    if not USE_HUB_STORAGE:
        if ON_SPACE:
            return (
                "Results were saved locally but not uploaded. "
                f"Add an HF_TOKEN secret with write access to {RESULTS_DATASET_REPO}."
            )
        return None

    filename = os.path.basename(local_path)
    try:
        _ensure_results_dataset()
        _api().upload_file(
            path_or_fileobj=local_path,
            path_in_repo=filename,
            repo_id=RESULTS_DATASET_REPO,
            repo_type="dataset",
            commit_message=f"Update leaderboard results for {model}",
            token=TOKEN,
        )
        logger.info("Uploaded %s to dataset %s", filename, RESULTS_DATASET_REPO)
        return None
    except Exception as exc:
        logger.exception("Could not upload %s to results dataset", filename)
        return f"Could not upload results to the dataset repo: {exc}"


def upload_submission_zip_to_hub(local_path: str, model: str) -> str | None:
    """Persist the full submission ZIP to the results dataset."""
    if not USE_HUB_STORAGE:
        if ON_SPACE:
            return (
                "Submission archive was saved locally but not uploaded. "
                f"Add an HF_TOKEN secret with write access to {RESULTS_DATASET_REPO}."
            )
        return None

    filename = os.path.basename(local_path)
    path_in_repo = f"{SUBMISSIONS_REPO_PREFIX}/{filename}"
    try:
        _ensure_results_dataset()
        _api().upload_file(
            path_or_fileobj=local_path,
            path_in_repo=path_in_repo,
            repo_id=RESULTS_DATASET_REPO,
            repo_type="dataset",
            commit_message=f"Update submission archive for {model}",
            token=TOKEN,
        )
        logger.info("Uploaded %s to dataset %s", path_in_repo, RESULTS_DATASET_REPO)
        return None
    except Exception as exc:
        logger.exception("Could not upload %s to results dataset", path_in_repo)
        return f"Could not upload submission archive to the dataset repo: {exc}"


def ensure_submission_zip_local(model: str) -> str | None:
    """Ensure the submission zip for *model* exists locally; download from Hub if needed."""
    from src.leaderboard.store import (
        model_to_zip_filename,
        submission_zip_path_for_model,
    )

    local_path = submission_zip_path_for_model(model)
    if os.path.isfile(local_path):
        return local_path

    if not USE_HUB_STORAGE:
        return None

    path_in_repo = f"{SUBMISSIONS_REPO_PREFIX}/{model_to_zip_filename(model)}"
    try:
        os.makedirs(SUBMISSIONS_DIR, exist_ok=True)
        _download_dataset_file(path_in_repo)
    except Exception as exc:
        logger.warning("Could not download submission archive %s: %s", path_in_repo, exc)
        return None

    return local_path if os.path.isfile(local_path) else None