storage.py

"""HF Hub bucket (or S3) access: download files and upload results.

Storage configuration is read from environment variables so the Space can be moved
between users / organisations without editing the code:

    FFASR_BUCKET_ID    HF dataset repo acting as the read/write bucket.
                       Default: ``whojavumusic/FFASR_Leaderboard-storage``.
    FFASR_RESULTS_PATH Path inside the bucket for the leaderboard CSV.
                       Default: ``results/leaderboard.csv``.
    FFASR_JOBS_PATH    Path inside the bucket for the job-state CSV.
                       Default: ``results/jobs_state.csv``.
    HF_TOKEN           Hub token with **write** access to ``FFASR_BUCKET_ID``.
                       ``ds_token`` is also accepted as a legacy fallback.
    token_for_ffasr_jobs  (or ``FFASR_JOBS_TOKEN``) Hub token used **only** to submit
                       and poll Hugging Face **Jobs** (billing account with Job credits).
                       ``HF_TOKEN`` is still passed into the job for bucket artifact upload.

For a private Space + private bucket, set ``HF_TOKEN`` as a Space repository secret
belonging to (or granted access to) the organisation that owns the bucket.
Set ``token_for_ffasr_jobs`` to the account/org that pays for Hub Jobs (not the Space
runtime default token).
"""

import io
import os
import shutil
import subprocess
import tempfile
from contextlib import nullcontext
from pathlib import Path

# --- Active: HuggingFace Bucket storage ---
STORAGE_BACKEND = "hf_bucket"
HF_BUCKET_ID = os.environ.get("FFASR_BUCKET_ID", "treble-technologies/FFASR_Leaderboard-storage")
DATASET_PREFIX = ""
RESULTS_PATH = os.environ.get("FFASR_RESULTS_PATH", "results/leaderboard.csv")
JOBS_STATE_PATH = os.environ.get("FFASR_JOBS_PATH", "results/jobs_state.csv")
# Append-only leaderboard history (one row per changed model per version). Used to
# reconstruct and display past leaderboard states from the version dropdown.
HISTORY_PATH = os.environ.get("FFASR_HISTORY_PATH", "results/leaderboard_history.csv")

HF_TOKEN = os.environ.get("HF_TOKEN") or os.environ.get("ds_token")

# Env names tried in order (HF Space secrets are often lowercase like ``token_for_ffasr_jobs``).
_JOBS_TOKEN_ENV_KEYS = (
    "token_for_ffasr_jobs",
    "TOKEN_FOR_FFASR_JOBS",
    "FFASR_JOBS_TOKEN",
    "HF_JOBS_TOKEN",
)


def token_for_ffasr_jobs() -> str | None:
    """Hub token for submitting/polling FFASR remote Jobs (billing account)."""
    for key in _JOBS_TOKEN_ENV_KEYS:
        tok = os.environ.get(key, "").strip()
        if tok:
            return tok
    return None


def require_token_for_ffasr_jobs() -> str:
    """Jobs billing token; never falls back to ``HF_TOKEN`` (avoids wrong namespace / 402)."""
    tok = token_for_ffasr_jobs()
    if tok:
        return tok
    keys = ", ".join(f"`{k}`" for k in _JOBS_TOKEN_ENV_KEYS)
    raise RuntimeError(
        f"Remote Hub Jobs require a dedicated jobs token: set one of {keys} as a Space secret "
        "(Hub token with Job credits and permission to run jobs). "
        "HF_TOKEN is only used for dataset bucket read/write on the Space and inside the job."
    )

try:
    from huggingface_hub import (
        HfApi,
        list_bucket_tree,
        download_bucket_files,
        batch_bucket_files,
    )

    hf_api = HfApi(token=HF_TOKEN)
except Exception:
    hf_api = None
    download_bucket_files = None
    batch_bucket_files = None
    list_bucket_tree = None

# --- Optional: AWS S3 (uncomment in init / wire here if needed) ---
# STORAGE_BACKEND = "s3"
# import boto3
# s3 = boto3.client("s3")
# DATASET_BUCKET = os.environ.get("DATASET_S3_BUCKET", "")
# RESULTS_BUCKET = os.environ.get("RESULTS_S3_BUCKET", DATASET_BUCKET)
# RESULTS_KEY = os.environ.get("RESULTS_S3_KEY", "results/leaderboard.csv")


def _ensure_xet_progress_reporter_compatible() -> None:
    """
    Some container images mix an older ``XetProgressReporter`` with a newer ``hf_api`` that
    calls ``notify_upload_complete`` between chunked uploads.
    """
    try:
        from huggingface_hub.utils._xet_progress_reporting import XetProgressReporter
    except ImportError:
        return
    if hasattr(XetProgressReporter, "notify_upload_complete"):
        return

    def notify_upload_complete(self) -> None:
        try:
            self._total_bytes_offset = getattr(self.data_processing_bar, "total", None) or 0
            self._total_transfer_bytes_offset = getattr(self.upload_bar, "total", None) or 0
        except Exception:
            pass

    XetProgressReporter.notify_upload_complete = notify_upload_complete  # type: ignore[method-assign]


def upload_to_bucket(
    bucket_id: str,
    *,
    add: list[tuple[str | Path | bytes, str]] | None = None,
    copy: list[tuple[str, str, str, str]] | None = None,
    delete: list[str] | None = None,
    token: str | bool | None = None,
) -> None:
    """
    Upload/copy/delete bucket files with progress bars disabled and Xet reporter compatibility.
    """
    if batch_bucket_files is None:
        # In-process hub lacks the bucket API (e.g. qwen-asr stack pins hub<1.0).
        # Fall back to an isolated uv env that has it.
        if not _bucket_uv_available():
            raise RuntimeError("huggingface_hub bucket support is not available")
        _upload_to_bucket_via_uv(
            bucket_id, add=add, copy=copy, delete=delete, token=token
        )
        return
    _ensure_xet_progress_reporter_compatible()
    try:
        from huggingface_hub.utils import disable_progress_bars
    except ImportError:
        ctx = nullcontext()
    else:
        ctx = disable_progress_bars()
    with ctx:
        batch_bucket_files(
            bucket_id,
            add=add,
            copy=copy,
            delete=delete,
            token=token,
        )


# huggingface_hub version that ships the bucket API, used by the isolated-subprocess
# fallback when the in-process hub is pinned <1.0 (e.g. the qwen-asr stack, whose
# transformers==4.57.x pin forces huggingface-hub<1.0 and so lacks bucket support).
_BUCKET_FALLBACK_HUB_SPEC = os.environ.get(
    "FFASR_BUCKET_FALLBACK_HUB_SPEC", "huggingface-hub>=1.14.0"
)

_BUCKET_DOWNLOAD_SNIPPET = """\
import os, sys
from huggingface_hub import download_bucket_files

bucket_id, remote_path, local_path = sys.argv[1], sys.argv[2], sys.argv[3]
download_bucket_files(
    bucket_id,
    files=[(remote_path, local_path)],
    token=os.environ.get("FFASR_BUCKET_DL_TOKEN") or None,
)
"""

_BUCKET_UPLOAD_SNIPPET = """\
import json, os, sys
from huggingface_hub import batch_bucket_files

try:
    from huggingface_hub.utils import disable_progress_bars
    disable_progress_bars()
except Exception:
    pass

manifest_path, bucket_id = sys.argv[1], sys.argv[2]
with open(manifest_path, encoding="utf-8") as f:
    manifest = json.load(f)
add = [(src, dst) for src, dst in manifest.get("add", [])] or None
copy = [tuple(c) for c in manifest.get("copy", [])] or None
delete = list(manifest.get("delete", [])) or None
batch_bucket_files(
    bucket_id,
    add=add,
    copy=copy,
    delete=delete,
    token=os.environ.get("FFASR_BUCKET_DL_TOKEN") or None,
)
"""


def _bucket_uv_available() -> bool:
    return shutil.which("uv") is not None


def _run_bucket_uv_snippet(
    snippet: str, args: list[str], *, token: str | None = None
) -> subprocess.CompletedProcess:
    """Run ``snippet`` in an isolated ``uv`` env that has the hub bucket API.

    Used when this process' ``huggingface_hub`` is pinned <1.0 (no bucket support).
    ``uv`` resolves an ephemeral, cached environment with a recent ``huggingface_hub``
    (plus ``hf_xet`` for Xet-backed buckets), independent of the job's pinned stack.
    The Hub token is passed via ``FFASR_BUCKET_DL_TOKEN`` (never on argv).
    """
    uv = shutil.which("uv") or "uv"
    cmd = [
        uv,
        "run",
        "--no-project",
        "--python",
        "3.12",
        "--with",
        _BUCKET_FALLBACK_HUB_SPEC,
        "--with",
        "hf_xet",
        "python",
        "-c",
        snippet,
        *args,
    ]
    env = dict(os.environ)
    tok = token or HF_TOKEN
    if tok:
        env["FFASR_BUCKET_DL_TOKEN"] = tok
    env.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")
    try:
        return subprocess.run(
            cmd,
            env=env,
            check=False,
            capture_output=True,
            text=True,
        )
    except FileNotFoundError as exc:
        raise RuntimeError(
            "huggingface_hub bucket support is not available in this environment and the "
            "`uv` fallback could not be launched (uv not found on PATH)."
        ) from exc


def _download_bucket_file_via_uv(path: str, local_path: str) -> None:
    proc = _run_bucket_uv_snippet(
        _BUCKET_DOWNLOAD_SNIPPET, [HF_BUCKET_ID, path, local_path]
    )
    if proc.returncode != 0 or not os.path.exists(local_path):
        tail = (proc.stderr or proc.stdout or "").strip()[-2000:]
        raise RuntimeError(
            "Bucket download via isolated uv environment failed "
            f"(exit {proc.returncode}) for '{path}'.\n{tail}"
        )


def _upload_to_bucket_via_uv(
    bucket_id: str,
    *,
    add: list[tuple[str | Path | bytes, str]] | None,
    copy: list[tuple[str, str, str, str]] | None,
    delete: list[str] | None,
    token: str | bool | None = None,
) -> None:
    """Upload/copy/delete bucket files via the isolated ``uv`` env (hub <1.0 fallback).

    ``bytes`` sources in ``add`` are materialized to temp files so the subprocess can
    pass plain file paths to ``batch_bucket_files``.
    """
    tmpdir = tempfile.mkdtemp(prefix="ffasr_bucket_up_")
    try:
        add_items: list[tuple[str, str]] = []
        for i, (src, dst) in enumerate(add or []):
            if isinstance(src, (bytes, bytearray)):
                src_path = os.path.join(tmpdir, f"add_{i}_{os.path.basename(dst) or 'blob'}")
                with open(src_path, "wb") as fh:
                    fh.write(src)
            else:
                src_path = os.fspath(src)
            add_items.append((src_path, dst))
        manifest = {
            "add": add_items,
            "copy": [list(c) for c in (copy or [])],
            "delete": list(delete or []),
        }
        import json as _json

        manifest_path = os.path.join(tmpdir, "manifest.json")
        with open(manifest_path, "w", encoding="utf-8") as fh:
            _json.dump(manifest, fh)
        snippet_token = token if isinstance(token, str) else None
        proc = _run_bucket_uv_snippet(
            _BUCKET_UPLOAD_SNIPPET, [manifest_path, bucket_id], token=snippet_token
        )
        if proc.returncode != 0:
            tail = (proc.stderr or proc.stdout or "").strip()[-2000:]
            raise RuntimeError(
                "Bucket upload via isolated uv environment failed "
                f"(exit {proc.returncode}).\n{tail}"
            )
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)


def download_bucket_file(path: str) -> str:
    """Download a single file from HF Bucket and return the local path.

    Uses the in-process bucket API when available; otherwise falls back to an isolated
    ``uv`` environment that has a hub version with bucket support (the qwen-asr stack
    pins huggingface-hub<1.0, which lacks the bucket API).
    """
    local_dir = tempfile.mkdtemp()
    local_path = os.path.join(local_dir, os.path.basename(path))
    if download_bucket_files is not None:
        download_bucket_files(
            HF_BUCKET_ID,
            files=[(path, local_path)],
            token=HF_TOKEN,
        )
    else:
        _download_bucket_file_via_uv(path, local_path)
    return local_path